Commit Test SRT

app.py

@@ -1,15 +1,18 @@
 import streamlit as st
 from transformers import pipeline
+from heapq import nlargest
 
-# Function to extract text from SRT-formatted text
 # Function to extract text from SRT-formatted text
 def extract_text_from_srt_text(srt_text):
     lines = srt_text.strip().split("\n\n")  # Split by empty lines to separate subtitles
     texts = [subtitle.split("\n")[2] for subtitle in lines if subtitle.strip()]  # Extract text from the third line of each subtitle
     return " ".join(texts)
 
-# Load summarization pipeline
-summarizer = pipeline("summarization")
+# Function to generate summary from text
+def generate_summary(text, summary_length):
+    summarizer = pipeline("summarization")
+    summary = summarizer(text, max_length=summary_length, min_length=30, do_sample=False)
+    return summary[0]["summary_text"]
 
 # Streamlit app
 st.title("SRT Summarization")
@@ -23,10 +26,16 @@ if st.button("Summarize"):
     if srt_text_input.strip():
         # Extract text from SRT-formatted text
         text_to_summarize = extract_text_from_srt_text(srt_text_input)
-        # Summarize text
-        summary = summarizer(text_to_summarize, max_length=1500, min_length=10, do_sample=False)
-        # Display summary
+        # Generate summary
+        summary = generate_summary(text_to_summarize, 150)  # You can adjust the summary length as needed
+        # Extract top 4 sentences
+        sentences = text_to_summarize.split(". ")
+        top_sentences = nlargest(4, sentences, key=len)
+        top_subjects = "\n".join(top_sentences)
+        # Display summary and top 4 subjects
         st.subheader("Summary:")
-        st.write(summary[0]["summary_text"])
+        st.write(summary)
+        st.subheader("Top 4 Subjects:")
+        st.write(top_subjects)
     else:
         st.warning("Please enter some SRT-formatted text.")