from smolagents import Tool
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
import torch


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_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}")

        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