Update README.md

README.md

@@ -41,20 +41,25 @@ Instruction Adherence: Retains high fidelity in understanding and following user
 
 **Loading the Model:**
 ```py
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Load the model and tokenizer
-model_name = "suayptalha/FastLlama-3.2-1B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Example usage
-input_text = "Solve for x: 2x + 3 = 7"
-inputs = tokenizer(input_text, return_tensors="pt")
-outputs = model.generate(**inputs)
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-print(response)
+import torch
+from transformers import pipeline
+
+model_id = "suayptalha/FastLlama-3.2-1B-Instruct"
+pipe = pipeline(
+    "text-generation",
+    model=model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+messages = [
+    {"role": "system", "content": "You are a friendly assistant named FastLlama."},
+    {"role": "user", "content": "Who are you?"},
+]
+outputs = pipe(
+    messages,
+    max_new_tokens=256,
+)
+print(outputs[0]["generated_text"][-1])
 ```
 
 **Dataset:**