Update model_usage.py

model_usage.py

@@ -1,18 +1,30 @@
-from transformers import AutoModelForCausalLM, GPT2Tokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 model = AutoModelForCausalLM.from_pretrained(
     "Girinath11/recursive-language-model-48m",
     trust_remote_code=True
 )
-tokenizer = GPT2Tokenizer.from_pretrained(
-    "Girinath11/recursive-language-model-48m"
-)
-print("Model loaded successfully!")
-prompt = "The future of artificial intelligence"
-input_ids = tokenizer.encode(prompt, return_tensors="pt")
-outputs = model.generate(
-    input_ids,
-    max_new_tokens=50,
-    temperature=0.8,
-    do_sample=True
-)
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+model.eval()
+print(f"Model loaded on {device}\n")
+prompts = [
+    "The future of artificial intelligence",
+    "Once upon a time",
+    "The key to success is"
+]
+for prompt in prompts:
+    print(f"Prompt: {prompt}")
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=50,
+            temperature=0.8,
+            top_p=0.9,
+            do_sample=True
+        )
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    print(f"{text}\n")