| import gradio as gr | |
| from transformers import GPT2LMHeadModel, GPT2Tokenizer | |
| import torch | |
| # Load model and tokenizer (using smaller GPT-2 for free tier) | |
| model_name = "gpt2" # You can also use "gpt2-medium" if it fits in memory | |
| tokenizer = GPT2Tokenizer.from_pretrained(model_name) | |
| model = GPT2LMHeadModel.from_pretrained(model_name) | |
| # Set pad token | |
| tokenizer.pad_token = tokenizer.eos_token | |
| def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9, top_k=50): | |
| """Generate text using GPT-2""" | |
| try: | |
| # Encode input | |
| inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True) | |
| # Generate | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| inputs, | |
| max_length=min(max_length + len(inputs[0]), 512), # Limit total length | |
| temperature=temperature, | |
| top_p=top_p, | |
| top_k=top_k, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id, | |
| num_return_sequences=1 | |
| ) | |
| # Decode output | |
| generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| # Return only the new generated part | |
| return generated_text[len(prompt):].strip() | |
| except Exception as e: | |
| return f"Error generating text: {str(e)}" | |
| # Create Gradio interface | |
| with gr.Blocks(title="GPT-2 Text Generator") as demo: | |
| gr.Markdown("# GPT-2 Text Generation Server") | |
| gr.Markdown("Enter a prompt and generate text using GPT-2. Free tier optimized!") | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt_input = gr.Textbox( | |
| label="Prompt", | |
| placeholder="Enter your text prompt here...", | |
| lines=3 | |
| ) | |
| with gr.Row(): | |
| max_length = gr.Slider( | |
| minimum=10, | |
| maximum=200, | |
| value=100, | |
| step=10, | |
| label="Max Length" | |
| ) | |
| 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.9, | |
| step=0.1, | |
| label="Top-p" | |
| ) | |
| top_k = gr.Slider( | |
| minimum=1, | |
| maximum=100, | |
| value=50, | |
| step=1, | |
| label="Top-k" | |
| ) | |
| generate_btn = gr.Button("Generate Text", variant="primary") | |
| with gr.Column(): | |
| output_text = gr.Textbox( | |
| label="Generated Text", | |
| lines=10, | |
| placeholder="Generated text will appear here..." | |
| ) | |
| # Examples | |
| gr.Examples( | |
| examples=[ | |
| ["Once upon a time in a distant galaxy,"], | |
| ["The future of artificial intelligence is"], | |
| ["In the heart of the ancient forest,"], | |
| ["The detective walked into the room and noticed"], | |
| ], | |
| inputs=prompt_input | |
| ) | |
| # Connect the function with explicit API endpoint name | |
| generate_btn.click( | |
| fn=generate_text, | |
| inputs=[prompt_input, max_length, temperature, top_p, top_k], | |
| outputs=output_text, | |
| api_name="/predict" # Explicit API endpoint for external calls | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| demo.launch() |