Update app.py

app.py

@@ -1,21 +1,29 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer
 
 # Load the tokenizer and model
 model_name = "JuliaUpton/Math_AI"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+mixtral_tokenizer = AutoTokenizer.from_pretrained(model_name)
+merged_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# 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 answer_question(question):
-    input_ids = tokenizer.encode(question, return_tensors="pt")
-    output_ids = model.generate(input_ids, max_length=100)
-    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return answer
+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="Answer")
+outputs = gr.outputs.Textbox(label="Response")
 
 examples = [
     ["What is the square root of 64?"],
@@ -25,11 +33,11 @@ examples = [
 
 # Launch the interface
 iface = gr.Interface(
-    fn=answer_question, 
+    fn=make_inference, 
     inputs=inputs,
     outputs=outputs,
-    title="Math Question Answering",
-    description="Enter a math question and get the answer with an explanation.",
+    title="Math Inquiry Answering",
+    description="Enter a math question and get a response with an aexplanation.",
     examples=examples
 )