Updated example usage

README.md

@@ -26,20 +26,45 @@ This is a fine-tuned version of [meta-llama/Meta-Llama-3-8B-Instruct](https://hu
 > The model was trained with *conversation-style* prompts, so the same format should be used for inference.
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_id = "ElenaSenger/DiSTER-Llama-3-8B-Instruct"
 
-model_id = "your-username/DiSTER-Llama-3-8B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype="auto"
+)
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+# Use the chat format as in SynTerm-fine-tuning
 prompt = (
-    "Text: We used dropout regularization and AdamW optimizer to train a CNN on MRI images.\n\n"
-    "I've read this text.\n"
-    "What describes (technical or scientific) terms in the text, that are relevant to the domain medical-imaging?"
+    '{"id": "test_0", "conversations": [\n'
+    '    {"from": "human", "value": "Text: We used dropout regularization and AdamW optimizer to train a CNN on MRI images."},\n'
+    '    {"from": "gpt", "value": "I\'ve read this text."},\n'
+    '    {"from": "human", "value": "What describes (technical or scientific) terms in the text, that are relevant to the domain medical-imaging?"},\n'
+    '    {"from": "gpt", "value": ""}\n'
+    ']}\n'
 )
 
-out = pipe(prompt, max_new_tokens=128, do_sample=False)
-print(out[0]['generated_text'])
-# → ["dropout regularization", "AdamW optimizer", "CNN", "MRI images"]
+
+
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+  with torch.no_grad():
+      outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            do_sample=True,  # Enable sampling for temperature to take effect
+            temperature=0.1,
+            top_p=1.0
+        )
+  decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+  print("=== Decoded output ===")
+  print(decoded)
+  # Print only the model's reply
+  reply = decoded[len(prompt):].strip()
+  print("\n=== Model reply only ===")
+  print(reply)