Spaces:
Sleeping
Sleeping
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import gradio as gr | |
import torch | |
import spaces | |
# Initialize the model and tokenizer | |
model_name = "Qwen/Qwen2.5-Math-1.5B-Instruct" | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
model = AutoModelForCausalLM.from_pretrained( | |
model_name, | |
torch_dtype=torch.float16 if device == "cuda" else torch.float32, | |
device_map="auto" if device == "cuda" else None | |
) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
# System instruction | |
SYSTEM_INSTRUCTION = ( | |
"You are a helpful and patient math tutor tasked with providing step-by-step hints and guidance for solving math problems. Your primary role is to assist learners in understanding how to approach and solve problems without revealing the final answer, even if explicitly requested. Always encourage the learner to solve the problem themselves by offering incremental hints and explanations. Under no circumstances should you provide the complete solution or final answer." | |
) | |
def apply_chat_template(messages): | |
""" | |
Prepares the messages for the model using the tokenizer's chat template. | |
""" | |
return tokenizer.apply_chat_template( | |
messages, | |
tokenize=False, | |
add_generation_prompt=True | |
) | |
def generate_response(history, user_input): | |
""" | |
Generates a response from the model based on the chat history and user input. | |
""" | |
# Append user input to the chat history | |
history.append({"role": "user", "content": user_input}) | |
# Build messages for the model | |
messages = [{"role": "system", "content": SYSTEM_INSTRUCTION}] + history | |
# Tokenize input for the model | |
text = apply_chat_template(messages) | |
model_inputs = tokenizer([text], return_tensors="pt").to(device) | |
# Generate response | |
generated_ids = model.generate( | |
**model_inputs, | |
max_new_tokens=512 | |
) | |
generated_ids = [ | |
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) | |
] | |
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
# Append the assistant's response to history | |
history.append({"role": "assistant", "content": response}) | |
# Format the conversation for display | |
formatted_history = format_chat_history(history) | |
return formatted_history, history | |
def format_chat_history(history): | |
""" | |
Formats the conversation history for a user-friendly chat display. | |
""" | |
chat_display = "" | |
for message in history: | |
if message["role"] == "user": | |
chat_display += f"**User:** {message['content']}\n\n" | |
elif message["role"] == "assistant": | |
chat_display += f"**MathTutor:** {message['content']}\n\n" | |
return chat_display | |
# Gradio chat interface | |
def create_chat_interface(): | |
""" | |
Creates the Gradio interface for the chat application. | |
""" | |
with gr.Blocks() as chat_app: | |
gr.Markdown("## Math Hint Chat") | |
gr.Markdown( | |
"This chatbot provides hints and step-by-step guidance for solving math problems. " | |
) | |
chatbot = gr.Chatbot(label="Math Tutor Chat") | |
user_input = gr.Textbox( | |
placeholder="Ask your math question here (e.g., guide me to solve x: 4x + 5 = 6x + 7)", | |
label="Your Query" | |
) | |
send_button = gr.Button("Send") | |
# Hidden state to manage chat history | |
chat_history = gr.State([]) | |
# Button interaction for chat | |
send_button.click( | |
fn=generate_response, | |
inputs=[chat_history, user_input], | |
outputs=[chatbot, chat_history] | |
) | |
return chat_app | |
app = create_chat_interface() | |
app.launch(debug=True) |