Update README.md

README.md

@@ -62,15 +62,27 @@ Users should validate outputs with certified medical professionals. This model i
 ## How to Get Started with the Model
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from transformers import pipeline
+
+model_id = "InferenceLab/MediLlama-3.2"
+pipe = pipeline(
+    "text-generation",
+    model=model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+messages = [
+    {"role": "system", "content": "You are a helpful Medical assistant."},
+    {"role": "user", "content": "Hi! How are you?"},
+]
+outputs = pipe(
+    messages,
+    max_new_tokens=256,
+)
+print(outputs[0]["generated_text"][-1])
 
-model = AutoModelForCausalLM.from_pretrained("InferenceLab/MediLlama-3.2")
-tokenizer = AutoTokenizer.from_pretrained("InferenceLab/MediLlama-3.2")
-
-input_text = "What are the symptoms of diabetes?"
-inputs = tokenizer(input_text, return_tensors="pt")
-outputs = model.generate(**inputs)
-print(tokenizer.decode(outputs[0]))
 ````
 
 ## Training Details