Commit Test SRT

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from transformers import pipeline
 from heapq import nlargest
+from keybert import KeyBERT
 
 # Function to extract text from SRT-formatted text
 def extract_text_from_srt_text(srt_text):
@@ -13,10 +14,13 @@ def generate_summary(text, summary_length):
     summarizer = pipeline("summarization")
     summary = summarizer(text, max_length=summary_length, min_length=30, do_sample=False)
     summary_text = summary[0]["summary_text"]
-    sentences = text.split(". ")
-    top_sentences = nlargest(4, sentences, key=len)
-    top_subjects = "\n".join(top_sentences)
-    return summary_text, top_subjects
+    return summary_text
+
+# Function to extract top 4 topics from text
+def extract_top_topics(text, n_top_topics):
+    model = KeyBERT('distilbert-base-nli-mean-tokens')
+    keywords = model.extract_keywords(text, keyphrase_ngram_range=(1, 3), stop_words='english', use_maxsum=True, nr_candidates=20, top_n=n_top_topics)
+    return keywords
 
 # Streamlit app
 st.title("SRT Summarization")
@@ -24,13 +28,8 @@ st.title("SRT Summarization")
 # Logo image URL
 logo_url = "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcQ6uQl0omK_PHXBbyaCHdmh3VjCo_Yvgwavmcs5XRF9Rkjx5FpflxyO4yfux6d2ojKsCOA&usqp=CAU"  # Replace with your logo image URL
 
-# Center the logo
-st.markdown(
-    f'<div style="display: flex; justify-content: center;">'
-    f'<img src="{logo_url}" style="width: 364px;">'
-    f'</div>',
-    unsafe_allow_html=True
-)
+# Display logo
+st.image(logo_url, width=364)
 
 # Text area for user to input SRT-formatted text
 srt_text_input = st.text_area("Paste SRT-formatted text here:")
@@ -43,12 +42,16 @@ if st.button("Summarize"):
         with st.spinner("Summarizing..."):
             # Extract text from SRT-formatted text
             text_to_summarize = extract_text_from_srt_text(srt_text_input)
-            # Generate summary and top subjects
-            summary, top_subjects = generate_summary(text_to_summarize, 150)  # You can adjust the summary length as needed
-            # Display summary and top subjects
+            # Generate summary
+            summary = generate_summary(text_to_summarize, 150)  # You can adjust the summary length as needed
+            # Display summary
             st.subheader("Summary:")
             st.write(summary)
-            st.subheader("Top 4 Subjects:")
-            st.write(top_subjects, bullet=True)  # Display as bullet points
+            # Extract top 4 topics
+            top_topics = extract_top_topics(text_to_summarize, 4)
+            # Display top 4 topics
+            st.subheader("Top 4 Topics:")
+            for topic in top_topics:
+                st.write(f"- {topic}")
     else:
         st.warning("Please enter some SRT-formatted text.")