Syrinx commited on
Commit
aa49290
1 Parent(s): b5cdfbc

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -0
app.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import GPT2Tokenizer, GPT2LMHeadModel
3
+
4
+ import streamlit as st
5
+
6
+
7
+ # Load the tokenizer and model
8
+ tokenizer = GPT2Tokenizer.from_pretrained('webtoon_tokenizer')
9
+ model = GPT2LMHeadModel.from_pretrained('webtoon_model')
10
+
11
+
12
+ # Define the app
13
+ def main():
14
+ st.title('Webtoon Description Generator')
15
+
16
+ # Get the input from the user
17
+ title = st.text_input('Enter the title of the Webtoon:', '')
18
+
19
+ # Generate the description
20
+ if st.button('Generate Description'):
21
+ with st.spinner('Generating...'):
22
+ description = generate_description(title)
23
+ st.success(description)
24
+
25
+
26
+ # Define the function that generates the description
27
+ def generate_description(title):
28
+ # Preprocess the input
29
+ input_text = f"{title.lower()}"
30
+ input_ids = tokenizer.encode(input_text, return_tensors='pt')
31
+
32
+ # Generate the output using the model
33
+ output = model.generate(
34
+ input_ids=input_ids,
35
+ max_length=256,
36
+ num_beams=4,
37
+ early_stopping=True,
38
+ no_repeat_ngram_size=2
39
+ )
40
+
41
+ # Convert the output to text
42
+ description = tokenizer.decode(output[0], skip_special_tokens=True)
43
+
44
+ return description
45
+
46
+
47
+ if __name__ == '__main__':
48
+ main()