Spaces:
Runtime error
Runtime error
| import torch | |
| from transformers import GPT2Tokenizer, GPT2LMHeadModel | |
| import streamlit as st | |
| # Load the tokenizer and model | |
| tokenizer = GPT2Tokenizer.from_pretrained('webtoon_tokenizer') | |
| model = GPT2LMHeadModel.from_pretrained('webtoon_model') | |
| # Define the app | |
| def main(): | |
| st.title('Webtoon Description Generator') | |
| # Get the input from the user | |
| title = st.text_input('Enter the title of the Webtoon:', '') | |
| # Generate the description | |
| if st.button('Generate Description'): | |
| with st.spinner('Generating...'): | |
| description = generate_description(title) | |
| st.success(description) | |
| # Check if GPU is available | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| else: | |
| device = torch.device("cpu") | |
| # Define the function that generates the description | |
| def generate_description(title): | |
| # Preprocess the input | |
| input_text = f"{title}" | |
| input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device) | |
| # Generate the output using the model | |
| output = model.generate( | |
| input_ids=input_ids, | |
| max_length=200, | |
| num_beams=4, | |
| early_stopping=True, | |
| no_repeat_ngram_size=2 | |
| ) | |
| # Convert the output to text | |
| description = tokenizer.decode(output[0], skip_special_tokens=True) | |
| return description | |
| if __name__ == '__main__': | |
| main() | |