Create Math_Wizard.py

Math_Wizard.py

@@ -0,0 +1,46 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-math-7b-instruct")
+model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-math-7b-instruct", torch_dtype=torch.bfloat16, device_map="cpu")
+
+
+from transformers import GenerationConfig
+model.generation_config = GenerationConfig.from_pretrained("deepseek-ai/deepseek-math-7b-instruct")
+model.generation_config.pad_token_id = model.generation_config.eos_token_id
+
+
+def solve_math_problem(questions):
+    if isinstance(questions, str):  # If input is a single string
+        questions = [questions]
+
+    results = []
+    for question in questions:
+        messages = [{"role": "user", "content": question}]
+        input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+        outputs = model.generate(input_tensor.to(model.device), max_new_tokens=100)
+        result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
+        results.append(result)
+    return results
+
+
+interface = gr.Interface(
+    fn=solve_math_problem,
+    inputs="text",
+    outputs="text",
+    title="Math Wizard",
+    description="""
+        Welcome to the Math Wizard!
+        Ask any math question, and let the wizard guide you through the solution step-by-step.
+    """,
+    allow_flagging=True,
+    examples=[
+        ["What is the integral of x^2?"],
+        ["How do I solve a quadratic equation?"]
+    ],
+    theme="compact"
+)
+
+interface.launch()