import streamlit as st
import torch
from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration

model_name = 'jian1114/jian_KoBART_title'
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name)
model = BartForConditionalGeneration.from_pretrained(model_name)

def process_paragraph(paragraph):
    # Return a list from tokenizer.encode instead of tensor
    input_ids_list = tokenizer.encode(paragraph, max_length=1024)
    # Convert the list to tensor when needed
    input_ids = torch.tensor([input_ids_list])
    
    output = model.generate(input_ids, max_length=32, num_beams=10, early_stopping=True)
    subheading = tokenizer.decode(output[0], skip_special_tokens=True)

    subheading_final = ""  # 실제 반환할 소제목
    check_list = ["em class", "violet_text", "green_text", "red_text", "blue_text"]
    if  subheading=="O" or "OO" in subheading:
        subheading_final = "😢소제목 생성 실패: 더 자세한 내용이 필요합니다." 
    elif any(x in subheading for x in check_list):
        subheading_final = "😢소제목 생성 실패: 문법 교정 후 다시 시도해 보세요."
    else:
        subheading_final = subheading

    return subheading_final


def main():
    css = """
    <style>
    textarea {
        height: 300px;
    }
    </style>
    """
    st.markdown(css, unsafe_allow_html=True)
    st.title("Subheading Generator")
    user_input = st.text_area("Enter a paragraph: ")

    if st.button("Generate"):
        if user_input:
            with st.spinner('Generating...'):
                result = process_paragraph(user_input)
                st.write(f'Subheading: {result}')
        else:
            st.warning('Please enter a paragraph.')

if __name__ == "__main__":
    main()