Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| import time | |
| import random | |
| # Model configuration - using TinyLlama for efficient CPU inference | |
| MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" | |
| # Global variables for model components | |
| tokenizer = None | |
| model = None | |
| text_generator = None | |
| def load_model(): | |
| """Load the Smol LLM model and tokenizer""" | |
| global tokenizer, model, text_generator | |
| try: | |
| print(f"Loading model: {MODEL_NAME}") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_NAME, | |
| torch_dtype=torch.float32, # Use float32 for CPU | |
| device_map="auto" | |
| ) | |
| # Create text generation pipeline | |
| text_generator = pipeline( | |
| "text-generation", | |
| model=model, | |
| tokenizer=tokenizer, | |
| max_new_tokens=512, | |
| temperature=0.7, | |
| top_p=0.95, | |
| do_sample=True | |
| ) | |
| # Set pad token if not present | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| return "✅ Model loaded successfully!" | |
| except Exception as e: | |
| return f"❌ Error loading model: {str(e)}" | |
| def format_prompt(prompt, system_prompt=None): | |
| """Format the prompt for chat-style models""" | |
| if system_prompt: | |
| formatted = f"<|system|>\n{system_prompt}\n<|user|>\n{prompt}\n<|assistant|>" | |
| else: | |
| formatted = f"<|user|>\n{prompt}\n<|assistant|>" | |
| return formatted | |
| def generate_text( | |
| prompt, | |
| max_length=200, | |
| temperature=0.7, | |
| top_p=0.95, | |
| repetition_penalty=1.1, | |
| system_prompt="You are a helpful AI assistant. Provide clear and concise answers." | |
| ): | |
| """Generate text using the loaded model""" | |
| global text_generator | |
| if text_generator is None: | |
| return "⚠️ Please load the model first using the 'Load Model' button." | |
| if not prompt.strip(): | |
| return "⚠️ Please enter a prompt." | |
| try: | |
| # Format the prompt | |
| formatted_prompt = format_prompt(prompt, system_prompt) | |
| # Update pipeline parameters | |
| text_generator.max_new_tokens = max_length | |
| text_generator.temperature = temperature | |
| text_generator.top_p = top_p | |
| text_generator.repetition_penalty = repetition_penalty | |
| # Generate response | |
| start_time = time.time() | |
| result = text_generator( | |
| formatted_prompt, | |
| max_new_tokens=max_length, | |
| temperature=temperature, | |
| top_p=top_p, | |
| repetition_penalty=repetition_penalty, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id, | |
| eos_token_id=tokenizer.eos_token_id | |
| ) | |
| generation_time = time.time() - start_time | |
| # Extract the generated text | |
| generated_text = result[0]["generated_text"] | |
| # Extract only the assistant's response | |
| if "<|assistant|>" in generated_text: | |
| response = generated_text.split("<|assistant|>")[-1].strip() | |
| else: | |
| response = generated_text | |
| # Format output with metadata | |
| output = f"**Response:**\n{response}\n\n---\n*Generated in {generation_time:.2f} seconds*" | |
| return output | |
| except Exception as e: | |
| return f"❌ Error during generation: {str(e)}" | |
| def clear_chat(): | |
| """Clear the chat interface""" | |
| return "", "" | |
| # Create custom theme | |
| custom_theme = gr.themes.Soft( | |
| primary_hue="blue", | |
| secondary_hue="indigo", | |
| neutral_hue="slate", | |
| font=gr.themes.GoogleFont("Inter"), | |
| text_size="lg", | |
| spacing_size="lg", | |
| radius_size="md" | |
| ).set( | |
| button_primary_background_fill="*primary_600", | |
| button_primary_background_fill_hover="*primary_700", | |
| block_title_text_weight="600", | |
| ) | |
| # Build the Gradio interface | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # 🤖 Smol LLM Inference GUI | |
| **Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)** - | |
| Efficient text generation using TinyLlama | |
| This application runs a compact language model locally for text generation. | |
| Perfect for chat, completion tasks, and creative writing. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| # Model loading section | |
| with gr.Group(): | |
| gr.Markdown("### 📦 Model Management") | |
| model_status = gr.Textbox( | |
| label="Model Status", | |
| value="Model not loaded. Click 'Load Model' to start.", | |
| interactive=False | |
| ) | |
| load_btn = gr.Button( | |
| "🔄 Load Model", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| # Generation parameters | |
| gr.Markdown("### ⚙️ Generation Parameters") | |
| with gr.Row(): | |
| max_length = gr.Slider( | |
| minimum=50, | |
| maximum=1024, | |
| value=200, | |
| step=50, | |
| label="Max Tokens" | |
| ) | |
| temperature = gr.Slider( | |
| minimum=0.1, | |
| maximum=2.0, | |
| value=0.7, | |
| step=0.1, | |
| label="Temperature" | |
| ) | |
| with gr.Row(): | |
| top_p = gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.95, | |
| step=0.05, | |
| label="Top-p" | |
| ) | |
| repetition_penalty = gr.Slider( | |
| minimum=1.0, | |
| maximum=2.0, | |
| value=1.1, | |
| step=0.1, | |
| label="Repetition Penalty" | |
| ) | |
| system_prompt = gr.Textbox( | |
| label="System Prompt", | |
| value="You are a helpful AI assistant. Provide clear and concise answers.", | |
| lines=3, | |
| placeholder="Enter a system prompt to guide the model's behavior..." | |
| ) | |
| with gr.Column(scale=3): | |
| # Main interface | |
| with gr.Group(): | |
| gr.Markdown("### 💬 Text Generation") | |
| prompt_input = gr.Textbox( | |
| label="Enter your prompt", | |
| placeholder="Type your message here...", | |
| lines=4, | |
| autofocus=True | |
| ) | |
| with gr.Row(): | |
| generate_btn = gr.Button( | |
| "🚀 Generate", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| clear_btn = gr.Button( | |
| "🗑️ Clear", | |
| variant="secondary" | |
| ) | |
| output_text = gr.Markdown( | |
| label="Generated Response", | |
| value="*Response will appear here...*" | |
| ) | |
| # Example prompts | |
| with gr.Accordion("📝 Example Prompts", open=False): | |
| gr.Examples( | |
| examples=[ | |
| ["Write a short story about a robot discovering music."], | |
| ["Explain quantum computing in simple terms."], | |
| ["Create a poem about the changing seasons."], | |
| ["What are the benefits of renewable energy?"], | |
| ["Write a Python function to calculate fibonacci numbers."], | |
| ["Describe the perfect day in your own words."], | |
| ["Explain the concept of machine learning to a beginner."], | |
| ["Create a dialogue between two friends planning a trip."] | |
| ], | |
| inputs=[prompt_input], | |
| label="Click an example to get started" | |
| ) | |
| # Event handlers | |
| load_btn.click( | |
| fn=load_model, | |
| outputs=[model_status], | |
| api_visibility="public" | |
| ) | |
| generate_btn.click( | |
| fn=generate_text, | |
| inputs=[ | |
| prompt_input, | |
| max_length, | |
| temperature, | |
| top_p, | |
| repetition_penalty, | |
| system_prompt | |
| ], | |
| outputs=[output_text], | |
| api_visibility="public" | |
| ) | |
| clear_btn.click( | |
| fn=clear_chat, | |
| outputs=[prompt_input], | |
| api_visibility="private" | |
| ) | |
| # Allow Enter key to generate | |
| prompt_input.submit( | |
| fn=generate_text, | |
| inputs=[ | |
| prompt_input, | |
| max_length, | |
| temperature, | |
| top_p, | |
| repetition_penalty, | |
| system_prompt | |
| ], | |
| outputs=[output_text], | |
| api_visibility="public" | |
| ) | |
| # Launch the application | |
| demo.launch( | |
| theme=custom_theme, | |
| footer_links=[ | |
| {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}, | |
| {"label": "TinyLlama Model", "url": "https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0"}, | |
| {"label": "Gradio", "url": "https://gradio.app"} | |
| ], | |
| share=False, | |
| show_error=True | |
| ) |