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