ZenithBrew commited on
Commit
0ad554f
·
verified ·
1 Parent(s): cdaf644

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -4
app.py CHANGED
@@ -1,8 +1,13 @@
1
  import gradio as gr
2
- from transformers import pipeline
3
 
4
- # Load the text generation model
5
- generator = pipeline('text-generation', model='gpt2')
 
 
 
 
 
6
 
7
  # Function to generate text
8
  def generate_text(prompt):
@@ -18,5 +23,5 @@ iface = gr.Interface(
18
  description="Generate text with AI!"
19
  )
20
 
21
- # Launch it now
22
  iface.launch()
 
1
  import gradio as gr
2
+ from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
3
 
4
+ # Pre-download the model and tokenizer
5
+ model_name = "gpt2"
6
+ GPT2Tokenizer.from_pretrained(model_name)
7
+ GPT2LMHeadModel.from_pretrained(model_name)
8
+
9
+ # Load the pipeline
10
+ generator = pipeline('text-generation', model=model_name)
11
 
12
  # Function to generate text
13
  def generate_text(prompt):
 
23
  description="Generate text with AI!"
24
  )
25
 
26
+ # Launch it with model!
27
  iface.launch()