Jiraheya commited on
Commit
51ac3df
1 Parent(s): 56c25b8

Add application and requirements file

Browse files
Files changed (2) hide show
  1. app.py +58 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, BartForConditionalGeneration
3
+
4
+ # Load the models and tokenizers
5
+ @st.cache_resource
6
+ def load_models():
7
+ t5_model = AutoModelForSeq2SeqLM.from_pretrained("Jiraheya/samsum_model_t5_small_10_epochs")
8
+ t5_tokenizer = AutoTokenizer.from_pretrained("Jiraheya/samsum_model_t5_small_10_epochs")
9
+
10
+ bart_model = BartForConditionalGeneration.from_pretrained("Jiraheya/pegasus_xsum_samsum_model_10epoch")
11
+ bart_tokenizer = AutoTokenizer.from_pretrained("Jiraheya/pegasus_xsum_samsum_model_10epoch")
12
+
13
+ return t5_model, t5_tokenizer, bart_model, bart_tokenizer
14
+
15
+ t5_model, t5_tokenizer, bart_model, bart_tokenizer = load_models()
16
+
17
+ # Set up the Streamlit app
18
+ st.title("Dialogue Summarizer Chatbot")
19
+
20
+ # Create a dropdown for model selection
21
+ model_choice = st.selectbox(
22
+ "Choose a model:",
23
+ ("T5-small", "BART-large-cnn")
24
+ )
25
+
26
+ # Create a text area for user input
27
+ user_input = st.text_area("Enter your dialogue here:", height=200)
28
+
29
+ # Create a button to generate summary
30
+ if st.button("Generate Summary"):
31
+ if user_input:
32
+ # Prepare input for the model
33
+ input_text = "summarize: " + user_input
34
+
35
+ if model_choice == "T5-small":
36
+ inputs = t5_tokenizer([input_text], max_length=1024, return_tensors="pt")
37
+ summary_ids = t5_model.generate(inputs["input_ids"], num_beams=2, min_length=0, max_length=60)
38
+ summary = t5_tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
39
+ else: # BART-large-cnn
40
+ inputs = bart_tokenizer([input_text], max_length=1024, return_tensors="pt")
41
+ summary_ids = bart_model.generate(inputs["input_ids"], num_beams=2, min_length=10, max_length=60)
42
+ summary = bart_tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
43
+
44
+ # Display the summary
45
+ st.subheader("Summary:")
46
+ st.write(summary)
47
+ else:
48
+ st.warning("Please enter some dialogue to summarize.")
49
+
50
+ # Add information about the app in the sidebar
51
+ st.sidebar.subheader("About the App")
52
+ st.sidebar.info(
53
+ "This app uses fine-tuned models to summarize dialogues. "
54
+ "Choose a model, enter your dialogue in the text area, and click 'Generate Summary' to get a concise summary."
55
+ )
56
+ st.sidebar.markdown("Models available:")
57
+ st.sidebar.markdown("- T5-small: Jiraheya/samsum_model_t5_small_10_epochs")
58
+ st.sidebar.markdown("- BART-large-cnn: Jiraheya/pegasus_xsum_samsum_model_10epoch")
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ transformers
2
+ torch
3
+ streamlit