jbraun19 commited on
Commit
f4e343b
·
1 Parent(s): 6a417ed

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -0
app.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import torch
3
+ from transformers import GPT2Tokenizer, GPT2LMHeadModel
4
+
5
+ # Load the GPT-2 tokenizer and model
6
+ tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
7
+ model = GPT2LMHeadModel.from_pretrained("gpt2")
8
+
9
+ # Set the maximum length of generated text
10
+ max_length = 200
11
+
12
+ # Define a function to generate text
13
+ def generate_text(prompt):
14
+ # Encode the prompt
15
+ input_ids = tokenizer.encode(prompt, return_tensors="pt")
16
+
17
+ # Generate text
18
+ output = model.generate(
19
+ input_ids=input_ids,
20
+ max_length=max_length,
21
+ num_beams=5,
22
+ no_repeat_ngram_size=2,
23
+ early_stopping=True
24
+ )
25
+
26
+ # Decode the generated text
27
+ text = tokenizer.decode(output[0], skip_special_tokens=True)
28
+
29
+ return text
30
+
31
+ # Set up the Streamlit app
32
+ st.title("GPT-2 Text Generator")
33
+
34
+ # Add a text input widget for the user to enter a prompt
35
+ prompt = st.text_input("Enter a prompt:")
36
+
37
+ # When the user clicks the "Generate" button, generate text
38
+ if st.button("Generate"):
39
+ with st.spinner("Generating text..."):
40
+ text = generate_text(prompt)
41
+ st.write(text)