Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| # Load pre-trained model and tokenizer | |
| model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B" | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| # Move model to GPU if available | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| def generate_response_stream(message: str, history: list) -> str: | |
| """ | |
| Generate a response in real-time using streaming. | |
| """ | |
| # Combine the conversation history with the new message | |
| full_input = "" | |
| for user_msg, bot_msg in history: | |
| full_input += f"User: {user_msg}\nBot: {bot_msg}\n" | |
| full_input += f"User: {message}\nBot: " | |
| # Tokenize input text | |
| inputs = tokenizer.encode(full_input, return_tensors="pt", max_length=512, truncation=True).to(device) | |
| # Generate response using the model with streaming | |
| past_key_values = None | |
| generated_tokens = [] | |
| for _ in range(100): # Max length of the response | |
| outputs = model( | |
| inputs, | |
| past_key_values=past_key_values, | |
| use_cache=True | |
| ) | |
| next_token_logits = outputs.logits[:, -1, :] | |
| next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1) | |
| # Decode the generated token | |
| generated_token = tokenizer.decode(next_token[0], skip_special_tokens=True) | |
| yield generated_token | |
| # Stop generation if the end-of-sequence token is encountered | |
| if next_token.item() == tokenizer.eos_token_id: | |
| break | |
| # Update inputs and past key values for the next step | |
| inputs = next_token | |
| past_key_values = outputs.past_key_values | |
| # Append the generated token to the output | |
| generated_tokens.append(generated_token) | |
| # Yield the final response | |
| final_response = "".join(generated_tokens) | |
| yield final_response | |
| # Create ChatInterface | |
| with gr.Blocks() as demo: | |
| gr.Markdown("""<center><font size=8>Chat with DeepSeek 🚀</center>""") | |
| chatbot = gr.Chatbot(label="Conversation") | |
| textbox = gr.Textbox(lines=2, label="Your Message") | |
| def add_user_message(message, history): | |
| """Add user message to the chat history.""" | |
| return "", history + [[message, None]] | |
| def bot_response(history): | |
| """Generate bot response using streaming.""" | |
| user_message = history[-1][0] | |
| history[-1][1] = "" | |
| for token in generate_response_stream(user_message, history[:-1]): | |
| history[-1][1] += token | |
| yield history | |
| # Event handlers | |
| submit_event = textbox.submit(add_user_message, [textbox, chatbot], [textbox, chatbot]).then( | |
| bot_response, chatbot, chatbot | |
| ) | |
| # Add a "Send" button | |
| send_button = gr.Button("🚀 Send") | |
| send_event = send_button.click(add_user_message, [textbox, chatbot], [textbox, chatbot]).then( | |
| bot_response, chatbot, chatbot | |
| ) | |
| # Launch the app | |
| demo.queue(api_open=False) | |
| demo.launch(max_threads=40) |