Update app.py

app.py

@@ -1,167 +1,160 @@
 import streamlit as st
-import ffmpeg
-import numpy as np
-import pandas as pd
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
-from transformers import pipeline
-import os
+import whisper
 import tempfile
+import os
 import chromadb
-
-# ✅ Chunked Upload Config (50 MB chunks to bypass 200 MB limit)
-MAX_CHUNK_SIZE = 50 * 1024 * 1024  # 50 MB
-
-def upload_file_in_chunks(file):
-    """Split file into 50 MB chunks to bypass the 200 MB limit."""
-    chunk_list = []
-    bytes_data = file.read()
-    total_size = len(bytes_data)
-    num_chunks = (total_size // MAX_CHUNK_SIZE) + 1
-
-    st.write(f"🔹 Splitting into {num_chunks} chunks...")
-
-    for i in range(num_chunks):
-        start = i * MAX_CHUNK_SIZE
-        end = min((i + 1) * MAX_CHUNK_SIZE, total_size)
-        chunk_list.append(bytes_data[start:end])
-
-    # Combine all chunks into a single CSV
-    full_csv_data = b"".join(chunk_list)
-    return full_csv_data
-
-# Initialize Chroma DB
-chroma_client = chromadb.Client()
-collection_name = "subtitle_chunks"
-
-# Hugging Face Whisper Model
-whisper_model = pipeline("automatic-speech-recognition", model="openai/whisper-base")
-
-# Function to extract audio from video
-def extract_audio(video_file, audio_file):
-    input_file = ffmpeg.input(video_file)
-    output_file = ffmpeg.output(input_file, audio_file, **{'vn': None, 'ar': 16000, 'ac': 1, 'f': 'wav'})
-    ffmpeg.run(output_file)
-
-# Function to transcribe audio using Whisper
-def transcribe_audio(audio_file):
-    result = whisper_model(audio_file)
-    return result['text']
-
-# Function to generate embeddings
-def generate_embedding(text):
-    model = SentenceTransformer('all-MiniLM-L6-v2')
-    return model.encode(text).tolist()
-
-# Function to split large CSV into smaller chunks
-def split_csv(csv_path, chunk_size=50000):
-    chunks = []
-    for i, chunk in enumerate(pd.read_csv(csv_path, chunksize=chunk_size)):
-        chunk_file = f"chunk_{i}.csv"
-        chunk.to_csv(chunk_file, index=False)
-        chunks.append(chunk_file)
-    return chunks
-
-# Store subtitle chunks in Chroma DB
-def store_chunks_in_chroma(chunk_files):
-    collection = chroma_client.create_collection(name=collection_name)
-
-    for chunk_file in chunk_files:
-        chunk = pd.read_csv(chunk_file)
-        
-        for idx, row in chunk.iterrows():
-            text = row['text']
-            embedding = generate_embedding(text)
-            
-            collection.add(
-                documents=[text],
-                metadatas=[{"source": chunk_file}],
-                ids=[f"{chunk_file}_{idx}"]
-            )
-
-# Search subtitles in Chroma DB
-def search_in_chroma(transcribed_text, top_k=10):
-    query_embedding = generate_embedding(transcribed_text)
-    collection = chroma_client.get_collection(name=collection_name)
-    
-    results = collection.query(
-        query_embeddings=[query_embedding],
-        n_results=top_k
-    )
-    
-    matches = []
-    for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
-        matches.append({"text": doc, "source": meta['source']})
+from pydub import AudioSegment
+from sentence_transformers import SentenceTransformer
+import numpy as np
+import pickle
+
+# 💡 Set Page Configuration
+st.set_page_config(page_title="Audio Subtitle Tool", layout="wide")
+
+# 🌐 Initialize ChromaDB
+db_client = chromadb.PersistentClient(path="./chroma_data")
+sub_db = db_client.get_or_create_collection(name="audio_subtitles")
+
+# ⚙️ Load Whisper Model
+@st.cache_resource()
+def load_transcriber():
+    return whisper.load_model("base")
+
+transcriber = load_transcriber()
+
+# ⚙️ Load Sentence Transformer
+@st.cache_resource()
+def load_encoder():
+    return SentenceTransformer("paraphrase-MiniLM-L6-v2")
+
+semantic_encoder = load_encoder()
+
+# 🔥 Load Subtitle Data from Pickle
+@st.cache_resource()
+def load_subtitle_store():
+    with open("subtitle_data.pkl", "rb") as file:
+        subtitle_store = pickle.load(file)
+
+# ⏱️ Format Time into SRT format
+def time_formatter(sec):
+    hr = int(sec // 3600)
+    mins = int((sec % 3600) // 60)
+    secs = int(sec % 60)
+    millis = int((sec % 1) * 1000)
+    return f"{hr:02}:{mins:02}:{secs:02},{millis:03}"
+
+# 🎙️ Transcribe Audio and Store in ChromaDB
+@st.cache_resource()
+def process_audio(file_path):
+    result = transcriber.transcribe(file_path, beam_size=1)
+    segments = result['segments']
+    srt_subtitles = []
+    plain_text_subs = []
+
+    for idx, seg in enumerate(segments):
+        start = time_formatter(seg['start'])
+        end = time_formatter(seg['end'])
+        content = seg['text']
+
+        # Create SRT format
+        srt_subtitles.append(f"{idx + 1}\n{start} --> {end}\n{content}\n")
+        plain_text_subs.append(content)
+
+        # Store in ChromaDB
+        sub_db.upsert(
+            documents=[content],
+            metadatas=[{"start_time": start, "end_time": end}],
+            ids=[f"subtitle_{idx}"]
+        )
+
+    return "\n".join(srt_subtitles), " ".join(plain_text_subs)
+
+# 🔍 Subtitle Search Function
+def search_content(query):
+    query_vec = semantic_encoder.encode(query).tolist()
+    matches = sub_db.query(query_embeddings=[query_vec], n_results=5)
+
+    if 'documents' in matches and matches['documents']:
+        results = []
+        for idx, doc in enumerate(matches['documents'][0]):
+            meta = matches['metadatas'][0][idx]
+            results.append(f"📌 {doc} (From: {meta['start_time']} → To: {meta['end_time']})")
+        return results
+    return []
+
+# 🔥 Sidebar Navigation
+st.sidebar.title("🎯 Navigation")
+page = st.sidebar.radio("Choose Action", ["Generate Subtitles", "Search Subtitles"])
+
+# 🎛️ Sidebar Styling
+st.sidebar.markdown(
+    """
+    <style>
+    .sidebar .sidebar-content {
+        background-color: #f0f0f5;
+        color: #333;
+    }
+    </style>
+    """,
+    unsafe_allow_html=True
+)
+
+# 🎬 Header Section
+st.markdown(
+    """
+    <h1 style='text-align: center; color: #673AB7;'>🎤 Audio Subtitle Generator & Search</h1>
+    <p style='text-align: center; font-size: 18px; color: #666;'>Generate subtitles from audio files and search through them.</p>
+    """,
+    unsafe_allow_html=True
+)
+
+# 🚀 Load Subtitle Data
+load_subtitle_store()
+
+# 🎵 File Upload Section
+st.markdown("### 🎧 Upload Your Audio File")
+audio_file = st.file_uploader("Supported formats: MP3, WAV", type=["mp3", "wav"])
+
+# 🎯 Page Logic
+if page == "Generate Subtitles":
+    if audio_file:
+        st.audio(audio_file, format='audio/mp3')
+
+        with st.spinner("🔧 Processing audio... Please wait."):
+            # Temporary file handling
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_audio:
+                if audio_file.name.endswith(".mp3"):
+                    audio = AudioSegment.from_mp3(audio_file)
+                elif audio_file.name.endswith(".wav"):
+                    audio = AudioSegment.from_wav(audio_file)
+                audio = audio.set_frame_rate(16000)
+                audio.export(tmp_audio.name, format="wav")
+
+                # Transcribe and store in ChromaDB
+                subtitle_output, plain_output = process_audio(tmp_audio.name)
+                st.success("✅ Transcription Completed!")
+
+                # Display Subtitles
+                st.markdown("### 📜 Generated Subtitles")
+                st.text_area("Subtitles (SRT Format)", subtitle_output, height=300)
+
+                # Download Options
+                st.download_button("⬇️ Download SRT File", subtitle_output, file_name="subtitles.srt", mime="text/plain")
+                st.download_button("⬇️ Download Plain Text", plain_output, file_name="subtitles.txt", mime="text/plain")
+
+                os.remove(tmp_audio.name)
+
+elif page == "Search Subtitles":
+    st.subheader("🔎 Search Subtitles")
+    query_input = st.text_input("Enter text to search")
     
-    return matches
-
-# Streamlit UI
-st.title("🎥 Video Subtitle Search with Hugging Face Whisper + Chroma DB")
-
-# Sidebar for CSV Upload with Chunking
-st.sidebar.header("📁 Upload Subtitle Database")
-csv_file = st.sidebar.file_uploader("Upload Subtitle CSV (200 MB+ supported)", type=["csv"])
-
-# Upload video
-uploaded_video = st.file_uploader("Upload a video", type=["mp4", "avi", "mov", "mkv"])
-
-# If both CSV and video are uploaded
-if csv_file and uploaded_video:
-    # Handle large CSV upload using 50 MB chunks
-    st.info("Uploading CSV in 50 MB chunks...")
-    full_csv_data = upload_file_in_chunks(csv_file)
-
-    # Save CSV locally
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_csv:
-        temp_csv.write(full_csv_data)
-        csv_path = temp_csv.name
-
-    # Save video locally
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
-        temp_video.write(uploaded_video.getbuffer())
-        video_path = temp_video.name
-
-    audio_path = "temp_audio.wav"
-
-    # Extract audio
-    st.info("Extracting audio...")
-    extract_audio(video_path, audio_path)
-
-    # Transcribe audio using Hugging Face Whisper
-    st.info("Transcribing audio...")
-    transcribed_text = transcribe_audio(audio_path)
-    st.text_area("Transcribed Text", transcribed_text, height=150)
-
-    # Split the CSV database into smaller chunks
-    st.info("Splitting CSV into smaller chunks...")
-    chunk_files = split_csv(csv_path, chunk_size=50000)
-
-    # Store chunks in Chroma DB
-    st.info("Storing subtitle chunks in Chroma DB...")
-    store_chunks_in_chroma(chunk_files)
-
-    # Search subtitles in Chroma DB
-    st.info("Searching subtitles in Chroma DB...")
-    matching_subtitles = search_in_chroma(transcribed_text)
-
-    # Display video
-    st.video(video_path)
-
-    # Display matching subtitles
-    st.subheader("📜 Matching Subtitles (Chroma DB + Cosine Similarity)")
-    for sub in matching_subtitles:
-        st.write(f"**Subtitle:** {sub['text']}")
-        st.write(f"**Source:** {sub['source']}")
-        st.write("---")
-
-    # Cleanup
-    os.remove(video_path)
-    os.remove(audio_path)
-    os.remove(csv_path)
-
-    # Remove chunk files
-    for chunk_file in chunk_files:
-        os.remove(chunk_file)
-
-else:
-    st.warning("Please upload both a video file and a CSV file to proceed.")
+    if query_input:
+        with st.spinner("🔍 Searching..."):
+            search_matches = search_content(query_input)
+            
+            if search_matches:
+                st.success("✅ Results Found:")
+                for match in search_matches:
+                    st.write(match)
+            else:
+                st.warning("⚠️ No matching results. Try a different query.")