Update myagent.py

myagent.py

@@ -11,9 +11,18 @@ import torch
 
 # Basic model wrapper for local inference with debug info
 class BasicAgent:
-    def __init__(self, model, tokenizer):
-
-        self.device = model.device if hasattr(model, 'device') else 'cpu'
+    def __init__(self):
+        print("[INFO] Loading default model...")
+        self.model = AutoModelForCausalLM.from_pretrained(
+            "LiquidAI/LFM2-1.2B",
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        )
+        
+        model_id="LiquidAI/LFM2-1.2B"
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.device = self.model.device if hasattr(self.model, 'device') else 'cpu'
         print(f"Model device: {self.device}")
 
     def _extract_prompt(self, prompt):
@@ -63,17 +72,6 @@ class BasicAgent:
     def __call__(self, prompt, max_new_tokens=512):
         return self.generate(prompt, max_new_tokens)
 
-# Load your model and tokenizer
-def load_model(model_id="LiquidAI/LFM2-1.2B"):
-    print(f"Loading model: {model_id}")
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-        trust_remote_code=True,
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    return BasicAgent(model, tokenizer)
 
 # Run minimal test
 if __name__ == "__main__":