Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import random | |
| import os | |
| # Force CPU usage if you're having GPU issues | |
| os.environ['CUDA_VISIBLE_DEVICES'] = '' | |
| # Load model and tokenizer | |
| model_name = "gpt2-medium" # You can change this to your fine-tuned model | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| # Set up padding token | |
| tokenizer.pad_token = tokenizer.eos_token | |
| model.config.pad_token_id = model.config.eos_token_id | |
| # Sample prompts for different styles | |
| style_prompts = { | |
| "Classic": "Here's a romantic pickup line: ", | |
| "Nerdy": "Here's a tech-related pickup line: ", | |
| "Flirty": "Here's a flirty pickup line: ", | |
| "Funny": "Here's a humorous pickup line: " | |
| } | |
| # Sample sentiment adjustments | |
| sentiment_adjustments = { | |
| "Sweet": " Make it sweet and genuine.", | |
| "Funny": " Make it lighthearted and humorous.", | |
| "Sarcastic": " Add a touch of playful sarcasm." | |
| } | |
| def generate_pickup_line(style, sentiment, temperature=0.7, max_length=50): | |
| """Generate a pickup line based on selected style and sentiment.""" | |
| try: | |
| prompt = style_prompts[style] + sentiment_adjustments[sentiment] | |
| # Set up the model inputs | |
| inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True) | |
| # Generate text | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| inputs["input_ids"], | |
| max_length=max_length, | |
| temperature=temperature, | |
| num_return_sequences=1, | |
| pad_token_id=tokenizer.pad_token_id, | |
| do_sample=True, | |
| no_repeat_ngram_size=2 | |
| ) | |
| # Decode and clean up the generated text | |
| generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| pickup_line = generated_text.replace(prompt, "").strip() | |
| return pickup_line if pickup_line else "Please try again!" | |
| except Exception as e: | |
| return f"An error occurred. Please try again! Error: {str(e)}" | |
| def copy_to_clipboard(text): | |
| """Copy the generated pickup line to clipboard.""" | |
| return text | |
| # Create the Gradio interface | |
| def create_interface(): | |
| with gr.Blocks(css="footer {visibility: hidden}") as interface: | |
| gr.Markdown("# 💘 AI Pickup Line Generator") | |
| gr.Markdown("Generate creative pickup lines with different styles and moods!") | |
| with gr.Row(): | |
| with gr.Column(): | |
| style = gr.Dropdown( | |
| choices=list(style_prompts.keys()), | |
| value="Classic", | |
| label="Pickup Line Style" | |
| ) | |
| sentiment = gr.Dropdown( | |
| choices=list(sentiment_adjustments.keys()), | |
| value="Sweet", | |
| label="Tone" | |
| ) | |
| generate_btn = gr.Button("Generate Pickup Line 💫", variant="primary") | |
| with gr.Column(): | |
| output = gr.Textbox( | |
| label="Your Generated Pickup Line", | |
| lines=3, | |
| interactive=False | |
| ) | |
| copy_btn = gr.Button("Copy to Clipboard 📋") | |
| # Handle generation | |
| generate_btn.click( | |
| fn=generate_pickup_line, | |
| inputs=[style, sentiment], | |
| outputs=output | |
| ) | |
| # Handle copying | |
| copy_btn.click( | |
| fn=copy_to_clipboard, | |
| inputs=output, | |
| outputs=None | |
| ) | |
| gr.Markdown(""" | |
| ### Tips: | |
| - Try different styles and tones to get varied results | |
| - Click 'Generate' multiple times to get different lines | |
| - Use the copy button to easily share your favorite lines | |
| """) | |
| return interface | |
| # Launch the interface | |
| if __name__ == "__main__": | |
| interface = create_interface() | |
| interface.launch() |