Spaces:
Paused
Paused
| import gradio as gr | |
| import torch | |
| from transformers import pipeline | |
| # Model configuration | |
| MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct" | |
| # Load the model once when the application starts | |
| device = 0 if torch.cuda.is_available() else -1 | |
| generator = pipeline( | |
| task="text-generation", | |
| model=MODEL_ID, | |
| device=device, | |
| torch_dtype="auto", | |
| ) | |
| def chat(message, history): | |
| """Generate an answer from the user's message.""" | |
| if not message or not message.strip(): | |
| return "Please enter a message." | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": "You are a helpful, concise AI assistant." | |
| } | |
| ] | |
| # Include previous conversation messages | |
| for user_message, assistant_message in history or []: | |
| messages.append({ | |
| "role": "user", | |
| "content": user_message | |
| }) | |
| messages.append({ | |
| "role": "assistant", | |
| "content": assistant_message | |
| }) | |
| messages.append({ | |
| "role": "user", | |
| "content": message.strip() | |
| }) | |
| try: | |
| result = generator( | |
| messages, | |
| max_new_tokens=256, | |
| do_sample=True, | |
| temperature=0.7, | |
| top_p=0.9, | |
| ) | |
| generated = result[0]["generated_text"] | |
| # Chat-template pipelines may return the full | |
| # conversation, so extract the last assistant reply. | |
| if isinstance(generated, list): | |
| answer = generated[-1]["content"] | |
| else: | |
| answer = generated | |
| return answer.strip() | |
| except Exception as error: | |
| print(f"Generation error: {error}") | |
| return "Sorry, I couldn't generate a response. Please try again." | |
| # Build the user interface | |
| demo = gr.ChatInterface( | |
| fn=chat, | |
| title="My Small AI Assistant", | |
| description=( | |
| "A lightweight AI chatbot powered by a " | |
| "pretrained Hugging Face language model." | |
| ), | |
| examples=[ | |
| "Explain Python in simple words.", | |
| "Write a short story about a robot.", | |
| "What is machine learning?", | |
| ], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |