Euryeth commited on
Commit
efdd63d
·
verified ·
1 Parent(s): 29cdb01

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -6
app.py CHANGED
@@ -1,9 +1,13 @@
1
  from transformers import pipeline
 
2
 
3
- # Initialize the text generation pipeline
4
- model = pipeline("text-generation", model="gpt2")
 
5
 
6
- def generate_text(prompt, max_length=100):
7
- """Generate text from a prompt"""
8
- output = model(prompt, max_length=max_length)
9
- return output[0]["generated_text"]
 
 
 
1
  from transformers import pipeline
2
+ import os
3
 
4
+ # Set cache directory to a writable location
5
+ os.environ['TRANSFORMERS_CACHE'] = '/tmp/cache'
6
+ os.environ['HF_HOME'] = '/tmp/cache'
7
 
8
+ # Initialize model with smaller config
9
+ model = pipeline(
10
+ "text-generation",
11
+ model="gpt2",
12
+ device=-1 # Force CPU usage
13
+ )