Update README.md

README.md

@@ -50,30 +50,66 @@ Then, copy the snippet from the section that is relevant for your usecase.
 
 ```python
 # pip install accelerate
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it")
+
+from .gemma2_inference_hf import get_completions
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerBase
+
+DTYPE = torch.bfloat16
+MODEL_ID = "idinsight/gemma-2-2b-it-ud"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, add_eos_token=False)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
 model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-2-2b-it",
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
+    MODEL_ID, device_map="auto", return_dict=True, torch_dtype=DTYPE
 )
-input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
-outputs = model.generate(**input_ids, max_new_tokens=32)
-print(tokenizer.decode(outputs[0]))
-```
 
-You can ensure the correct chat template is applied by using `tokenizer.apply_chat_template` as follows:
-```python
-messages = [
-    {"role": "user", "content": "Write me a poem about Machine Learning."},
-]
-input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True).to("cuda")
-outputs = model.generate(**input_ids, max_new_tokens=256)
-print(tokenizer.decode(outputs[0]))
+text_generation_params = {
+    "do_sample": True,
+    "eos_token_id": tokenizer.eos_token_id,
+    "max_new_tokens": 1024,
+    "num_return_sequences": 1,
+    "repetition_penalty": 1.1,
+    "temperature": 1e-6,
+    "top_p": 0.9,
+}
+
+response = get_completions(
+    model=model,
+    rules_list=[
+        "NOT URGENT",
+        "Bleeding from the vagina",
+        "Bad tummy pain",
+        "Bad headache that won’t go away",
+        "Bad headache that won’t go away",
+        "Changes to vision",
+        "Trouble breathing",
+        "Hot or very cold, and very weak",
+        "Fits or uncontrolled shaking",
+        "Baby moves less",
+        "Fluid from the vagina",
+        "Feeding problems",
+        "Fits or uncontrolled shaking",
+        "Fast, slow or difficult breathing",
+        "Too hot or cold",
+        "Baby’s colour changes",
+        "Vomiting and watery poo",
+        "Infected belly button",
+        "Swollen or infected eyes",
+        "Bulging or sunken soft spot",
+    ],
+    skip_special_tokens_during_decode=False,
+    text_generation_params=text_generation_params,
+    tokenizer=tokenizer,
+    user_message="If my newborn can't able to breathe what can i do",
+)
+print(f"{response = }")
 ```
 
+The `gemma2_inferece_hf.py` module is provided for downloaded with the model files.
+
 ### Evaluation
 
 XXX