Update app.py

app.py

@@ -2,12 +2,20 @@ 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: ",
@@ -25,27 +33,31 @@ sentiment_adjustments = {
 
 def generate_pickup_line(style, sentiment, temperature=0.7, max_length=50):
     """Generate a pickup line based on selected style and sentiment."""
-    prompt = style_prompts[style] + sentiment_adjustments[sentiment]
-    
-    # Set up the model inputs
-    inputs = tokenizer(prompt, return_tensors="pt", padding=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.eos_token_id,
-            do_sample=True
-        )
-    
-    # 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
+    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."""