Spaces:
Sleeping
Sleeping
File size: 4,046 Bytes
2a8ae3c 4a18f0c 2a8ae3c 4a18f0c 2a8ae3c 4a18f0c 2a8ae3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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() |