| from smolagents import Tool | |
| import hashlib | |
| import json | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import os | |
| class ModelMathTool(Tool): | |
| name = "math_model" | |
| description = "Answers advanced math questions using a pretrained math model." | |
| inputs = { | |
| "problem": { | |
| "type": "string", | |
| "description": "Math problem to solve.", | |
| } | |
| } | |
| 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 forward(self, problem: str) -> str: | |
| print(f"[MathModelTool] Question: {problem}") | |
| response = self.model.__call__(problem) | |
| return response | |