|
import torch |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
import gradio as gr |
|
|
|
|
|
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): |
|
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=False, |
|
examples=[ |
|
["What is the integral of x^2?"], |
|
["How do I solve a quadratic equation?"], |
|
["Tell me about Ramanujan"] |
|
] |
|
) |
|
|
|
if __name__ == "__main__": |
|
interface.launch() |
|
|