Update app.py

app.py

@@ -1,131 +1,128 @@
-import os
 import streamlit as st
+import ffmpeg
+import numpy as np
 import pandas as pd
-import chromadb
-from langchain.embeddings import HuggingFaceEmbeddings
-from moviepy.editor import VideoFileClip
-import whisper
-import tempfile
-import time
+from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
-import numpy as np
+from transformers import pipeline
+import os
+import tempfile
+import chromadb
 
-# Set up ChromaDB client
+# Chroma DB initialization
 chroma_client = chromadb.Client()
-collection_name = "subtitle_collection"
-
-if collection_name not in [c.name for c in chroma_client.list_collections()]:
+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 Hugging Face 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
+
+# Function to store subtitle chunks in Chroma DB
+def store_chunks_in_chroma(chunk_files):
     collection = chroma_client.create_collection(name=collection_name)
-else:
-    collection = chroma_client.get_collection(name=collection_name)
-
-# Load Whisper model
-st.title("🎥 AI-Powered Subtitle Extraction with Chunked Uploading")
-model = whisper.load_model("base")
-
-# ✅ 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
-
-# ✅ Upload CSV file (chunked)
-st.sidebar.header("📥 Upload CSV File (Subtitles DB)")
-uploaded_csv = st.sidebar.file_uploader("Upload large CSV file (up to 1 GB)", type=["csv"])
 
-if uploaded_csv:
-    st.sidebar.success("✅ CSV uploaded successfully!")
+    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}"]
+            )
+
+# Function to search subtitles using 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)
     
-    # Read file in chunks
-    full_csv_data = upload_file_in_chunks(uploaded_csv)
+    results = collection.query(
+        query_embeddings=[query_embedding],
+        n_results=top_k
+    )
     
-    # Save combined file
-    with open("uploaded_large_file.csv", "wb") as f:
-        f.write(full_csv_data)
-
-    # Load CSV into DataFrame
-    df = pd.read_csv("uploaded_large_file.csv")
-    st.sidebar.write("📊 Preview of CSV data:")
-    st.sidebar.dataframe(df.head())
-
-    # ✅ Generate embeddings and store in ChromaDB
-    st.sidebar.write("🔹 Generating embeddings and storing in ChromaDB...")
-
-    embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-
-    subtitles = df["subtitle"].tolist()
-
-    for i, subtitle in enumerate(subtitles):
-        emb = embeddings_model.embed_query(subtitle)  # Generate embedding
-        collection.add(
-            ids=[f"subtitle_{i}"],
-            embeddings=[emb],
-            metadatas=[{"text": subtitle}]
-        )
-
-    st.sidebar.success("✅ Subtitles stored in ChromaDB!")
+    matches = []
+    for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
+        matches.append({"text": doc, "source": meta['source']})
+    
+    return matches
 
-# ✅ Upload Video file
-st.sidebar.header("🎥 Upload Video File")
-uploaded_video = st.sidebar.file_uploader("Upload a video", type=["mp4", "avi", "mov"])
+# Streamlit UI
+st.title("🎥 Video Subtitle Search with Hugging Face Whisper + Chroma DB")
 
-if uploaded_video:
-    st.sidebar.success("✅ Video uploaded successfully!")
+# Upload video
+uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov", "mkv"])
 
-    # Save the video temporarily
+if uploaded_file:
     with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
-        temp_video.write(uploaded_video.read())
+        temp_video.write(uploaded_file.getbuffer())
         video_path = temp_video.name
 
-    st.sidebar.write("🔹 Extracting audio and transcribing...")
-    start_time = time.time()
+    audio_path = "temp_audio.wav"
 
-    # Extract audio and transcribe using Whisper
-    clip = VideoFileClip(video_path)
-    audio_path = f"{video_path}.wav"
-    clip.audio.write_audiofile(audio_path)
+    # Extract audio
+    st.info("Extracting audio...")
+    extract_audio(video_path, audio_path)
 
-    result = model.transcribe(audio_path)
-    transcribed_text = result["text"]
+    # 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)
 
-    st.write("📝 Transcribed Text:")
-    st.write(transcribed_text)
+    # Split the large CSV database into smaller chunks
+    st.info("Splitting large database into smaller chunks...")
+    subtitle_db_path = "database.csv"
+    chunk_files = split_csv(subtitle_db_path, chunk_size=50000)
 
-    # ✅ Retrieve matching subtitles from ChromaDB using Cosine Similarity
-    st.write("🔍 Retrieving matching subtitles...")
-    
-    # Embed the transcribed text
-    transcribed_emb = embeddings_model.embed_query(transcribed_text)
+    # Store chunks in Chroma DB
+    st.info("Storing subtitle chunks in Chroma DB...")
+    store_chunks_in_chroma(chunk_files)
 
-    # Query ChromaDB
-    results = collection.query(
-        query_embeddings=[transcribed_emb],
-        n_results=5  # Top 5 similar subtitles
-    )
+    # Search subtitles in Chroma DB
+    st.info("Searching subtitles in Chroma DB...")
+    matching_subtitles = search_in_chroma(transcribed_text)
 
-    # Display matching subtitles
-    st.write("✅ Matching Subtitles:")
-    for res in results["metadatas"][0]:
-        st.write(f"- {res['text']}")
+    # Display video
+    st.video(video_path)
 
-    end_time = time.time()
-    st.write(f"⏱️ Process completed in {end_time - start_time:.2f} seconds.")
+    # 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("---")
 
-    # Clean up temporary files
+    # Cleanup
     os.remove(video_path)
     os.remove(audio_path)
+
+    # Remove chunk files
+    for chunk_file in chunk_files:
+        os.remove(chunk_file)