import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from peft import PeftModel

# Load the tokenizer and model
model_name = "JuliaUpton/Math_AI"
model_architecture = "mistralai/Mixtral-8x7B-Instruct-v0.1"
mixtral_tokenizer = AutoTokenizer.from_pretrained(model_architecture)
mixtral_config = AutoConfig.from_pretrained(model_architecture, trust_remote_code=True)
merged_model = PeftModel.from_pretrained(model_name, model=AutoModelForCausalLM.from_pretrained(model_architecture, config=mixtral_config, trust_remote_code=True, device_map="auto"))

# Define the input formatting function
def input_from_text(instruction):
    return f"<s>[INST]Below is a math inquiry, please answer it as a math expert showing your thought process.\n\n### Inquiry:\n{instruction}\n\n### Response:[/INST]"

# Define the inference function
def make_inference(instruction):
    inputs = mixtral_tokenizer(input_from_text(instruction), return_tensors="pt")
    outputs = merged_model.generate(
        **inputs,
        max_new_tokens=150,
        generation_kwargs={"repetition_penalty": 1.7}
    )
    result = mixtral_tokenizer.decode(outputs[0], skip_special_tokens=True).split("[/INST]")[1]
    return result

# Create a Gradio interface
inputs = gr.inputs.Textbox(label="Enter your math question")
outputs = gr.outputs.Textbox(label="Response")

examples = [
    ["What is the square root of 64?"],
    ["Simplify: (3x^2 + 4x - 7) - (2x^2 - 3x + 5)"],
    ["Find the derivative of f(x) = 3x^3 - 2x^2 + 5x - 1"]
]

# Launch the interface
iface = gr.Interface(
    fn=make_inference, 
    inputs=inputs,
    outputs=outputs,
    title="Math Inquiry Answering",
    description="Enter a math question and get a response with an explanation.",
    examples=examples
)

iface.launch()