Update model_usage.py

model_usage.py

@@ -1,9 +1,6 @@
-from mixture_of_recursion import RecursiveLanguageModel, RecursiveLanguageModelConfig
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-
-# Load model
-model = RecursiveLanguageModel.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     "Girinath11/recursive-language-model-198m",
     trust_remote_code=True
 )
@@ -11,25 +8,18 @@ tokenizer = AutoTokenizer.from_pretrained(
     "Girinath11/recursive-language-model-198m",
     trust_remote_code=True
 )
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = model.to(device)
 model.eval()
-
 print(f"✅ Model loaded on {device}")
 print(f"📊 Parameters: {sum(p.numel() for p in model.parameters()):,}\n")
-
-
 def chat(question, max_new_tokens=150, temperature=0.7, top_p=0.9):
-    # Must use chat format — model trained on this
     prompt = f"<|user|>\n{question}\n<|assistant|>\n"
-
     inputs = tokenizer(
         prompt,
         return_tensors="pt",
         add_special_tokens=False
     ).to(device)
-
     with torch.no_grad():
         outputs = model.generate(
             inputs['input_ids'],
@@ -38,31 +28,19 @@ def chat(question, max_new_tokens=150, temperature=0.7, top_p=0.9):
             top_p=top_p,
             do_sample=True,
         )
-
     full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Extract only assistant response
     if "<|assistant|>" in full_text:
         response = full_text.split("<|assistant|>")[-1].strip()
     else:
         response = full_text.replace(question, "").strip()
 
     return response
-
-
-# Test
 questions = [
     "What is machine learning?",
     "What is Python programming?",
     "Explain neural networks simply",
     "What is artificial intelligence?",
 ]
-
-print("=" * 55)
-print("🤖 MIXTURE OF RECURSION LM — 198M")
-print("=" * 55)
-
 for q in questions:
     print(f"\n❓ {q}")
-    print(f"💬 {chat(q)}")
-    print("-" * 55)
+    print(f"💬 {chat(q)}")