Update app.py

app.py

@@ -1,13 +1,20 @@
-from transformers import pipeline
 import os
+from transformers import pipeline
 
-# Set cache directory to a writable location
+# Force cache location BEFORE importing transformers
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/cache'
 os.environ['HF_HOME'] = '/tmp/cache'
+os.environ['HF_DATASETS_CACHE'] = '/tmp/cache'
+os.environ['HUGGINGFACE_HUB_CACHE'] = '/tmp/cache'
 
-# Initialize model with smaller config
+# Now import pipeline
 model = pipeline(
-    "text-generation",
-    model="gpt2",
-    device=-1  # Force CPU usage
-)
+    "text-generation", 
+    model="distilgpt2",  # Smaller model for reliability
+    device=-1  # Force CPU
+)
+
+def generate_text(prompt, max_length=100):
+    """Generate text from a prompt"""
+    output = model(prompt, max_length=max_length)
+    return output[0]["generated_text"]