updat4e application file

app.py

@@ -1,4 +1,38 @@
 import streamlit as st
+import torch
+from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+model_name = 'jian1114/jian_KoBART_subheading'
+tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name)
+model = BartForConditionalGeneration.from_pretrained(model_name)
+
+def process_paragraph(paragraph):
+    input_ids = tokenizer.encode(paragraph, return_tensors='pt', max_length=1024)
+    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():
+    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()