app.py

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)


# Define the function that generates the description
def generate_description(title):
    # Preprocess the input
    input_text = f"{title.lower()}"
    input_ids = tokenizer.encode(input_text, return_tensors='pt')

    # Generate the output using the model
    output = model.generate(
        input_ids=input_ids,
        max_length=256,
        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()