Update README.md

README.md

@@ -32,43 +32,39 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_id = "omdeep22/Gonyai-v1"
 
-# 1. Load Tokenizer and Model
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id, 
     trust_remote_code=True,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to("cuda" if torch.cuda.is_available() else "cpu")
-
-# 2. Define your prompt using the Chat Template
-# This automatically handles the <|user|> and <|assistant|> tags
-messages = [
-    {"role": "user", "content": "गोंयच्या पावसाचेर एक कविता बरोव."}
-]
+    torch_dtype=dtype
+).to(device)
 
-tokenized_chat = tokenizer.apply_chat_template(
+messages = [{"role": "user", "content": "गोंयच्या पावसाचेर एक कविता बरोव."}]
+inputs = tokenizer.apply_chat_template(
     messages, 
     tokenize=True, 
     add_generation_prompt=True, 
-    return_tensors="pt"
-).to(model.device)
+    return_tensors="pt",
+    return_dict=True
+).to(device)
 
-# 3. Optimized Inference Settings for 160M Architecture
-outputs = model.generate(
-    tokenized_chat,
-    max_new_tokens=80,          # Prevents rambling/hallucinations
-    min_new_tokens=10,          # Ensures a meaningful response
-    temperature=0.3,            # Keeps the model focused and logical
-    top_k=40,                   # Filters out low-probability noise
-    top_p=0.9,                  # Nucleus sampling for coherence
-    repetition_penalty=1.2,     # Prevents looping in small models
-    do_sample=True,
-    eos_token_id=tokenizer.eos_token_id,
-    pad_token_id=tokenizer.eos_token_id
-)
+with torch.inference_mode():
+    with torch.autocast(device_type=device, dtype=dtype):
+        outputs = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            max_new_tokens=100,
+            temperature=0.3,
+            repetition_penalty=1.2,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
-# 4. Decode only the NEW tokens generated by the assistant
-generated_tokens = outputs[0][tokenized_chat.shape[-1]:]
-response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
+response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
-print(f"Assistant: {response}")
+print(f"\nAssistant: {response}")