Commit Test SRT

app.py

@@ -1,7 +1,10 @@
 import streamlit as st
+from flask import Flask, request, jsonify
 from transformers import pipeline
 from keybert import KeyBERT
 
+# Initialize Flask app
+app = Flask(__name__)
 
 # Function to extract text from SRT-formatted text
 def extract_text_from_srt_text(srt_text):
@@ -9,7 +12,6 @@ def extract_text_from_srt_text(srt_text):
     texts = [subtitle.split("\n")[2] for subtitle in lines if subtitle.strip()]
     return " ".join(texts)
 
-
 # Function to generate summary from text
 def generate_summary(text, summary_length):
     summarizer = pipeline("summarization")
@@ -17,43 +19,48 @@ def generate_summary(text, summary_length):
     summary_text = summary[0]["summary_text"]
     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
-
+    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 [topic for topic, _ in keywords]
 
 # Streamlit app
 st.title("SRT Summarization")
 
+# Text area for user to input SRT-formatted text
+srt_text_input = st.text_area("Paste SRT-formatted text here:")
+
+# Button to trigger summarization
+if st.button("Summarize"):
+    # Check if text area is not empty
+    if srt_text_input.strip():
+        # Extract text from SRT-formatted text
+        text_to_summarize = extract_text_from_srt_text(srt_text_input)
+        # Generate summary
+        summary = generate_summary(text_to_summarize, 150)
+        # Extract top 4 topics
+        top_topics = extract_top_topics(text_to_summarize, 4)
+        # Display summary and top 4 topics
+        st.subheader("Summary:")
+        st.write(summary)
+        st.subheader("Top 4 Keywords:")
+        for topic in top_topics:
+            st.write(f"- {topic}")
+    else:
+        st.warning("Please enter some SRT-formatted text.")
 
-# Define summarization route
-@st.experimental_memo
-def summarize_text(srt_text):
+# Define endpoint for REST API
+@app.route("/summarize", methods=["POST"])
+def summarize():
+    data = request.json
+    if "srt_text" not in data:
+        return jsonify({"error": "Missing 'srt_text' parameter"}), 400
+    srt_text = data["srt_text"]
     text_to_summarize = extract_text_from_srt_text(srt_text)
     summary = generate_summary(text_to_summarize, 150)
     top_topics = extract_top_topics(text_to_summarize, 4)
-    return summary, top_topics
-
-
-# Main entry point
-def main():
-    srt_text_input = st.text_area("Paste SRT-formatted text here:")
-
-    if st.button("Summarize"):
-        if srt_text_input.strip():
-            summary, top_topics = summarize_text(srt_text_input)
-            st.subheader("Summary:")
-            st.write(summary)
-            st.subheader("Top 4 Keywords:")
-            for topic, _ in top_topics:
-                st.write(f"- {topic}")
-        else:
-            st.warning("Please enter some SRT-formatted text.")
-
+    return jsonify({"summary": summary, "top_topics": top_topics})
 
 if __name__ == "__main__":
-    main()
+    app.run()