Update tooling.py

tooling.py

@@ -1,8 +1,6 @@
-from smolagents import Tool, HfApiModel
-import hashlib
-import json
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
+from smolagents import Tool
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+import torch
 
 
 class ModelMathTool(Tool):
@@ -18,13 +16,32 @@ class ModelMathTool(Tool):
 
     output_type = "string"
 
-    def __init__(self, model_id="Qwen/Qwen2.5-Math-7B"):
-        print(f"Loading math model: {model_id}")
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-        self.model = HfApiModel(model_id=model_id, max_tokens=512)
+    def __init__(self, model_name= "deepseek-ai/deepseek-math-7b-base"):
+        print(f"Loading math model: {model_name}")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        print("loaded tokenizer")
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+        print("loaded auto model")
+
+        self.model.generation_config = GenerationConfig.from_pretrained(model_name)
+        print("loaded coonfig")
+
+        self.model.generation_config.pad_token_id = self.model.generation_config.eos_token_id
+        print("loaded pad token")
+
+
 
     def forward(self, problem: str) -> str:
         print(f"[MathModelTool] Question: {problem}")
-        response = self.model.__call__(problem)
-        return response
+
+        inputs = self.tokenizer(problem, return_tensors="pt")
+        outputs =self.model.generate(**inputs, max_new_tokens=100)
+
+        result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        return result
+
+
+