Spaces:
Sleeping
Sleeping
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): | |
lines = srt_text.strip().split("\n\n") | |
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") | |
summary = summarizer(text, max_length=summary_length, min_length=30, do_sample=False) | |
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 [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 endpoint for REST API | |
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 jsonify({"summary": summary, "top_topics": top_topics}) | |
if __name__ == "__main__": | |
app.run() | |