Update app.py

app.py

@@ -7,9 +7,61 @@ from sklearn.metrics.pairwise import cosine_similarity
 from transformers import pipeline
 import os
 import tempfile
-import math
-
-# Hugging Face Whisper Model
+import shutil
+import chromadb
+
+# Initialize Chroma DB client
+client = chromadb.Client()
+
+# Sidebar for CSV Upload and Permanent Save
+st.sidebar.title("📂 Upload CSV File")
+csv_file = st.sidebar.file_uploader("Choose a CSV file", type=["csv"])
+
+# Save CSV permanently
+def save_csv_permanently(uploaded_file):
+    save_path = os.path.join(os.getcwd(), "permanent_subtitle_data.csv")
+    with open(save_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+    return save_path
+
+# Load the CSV into Chroma DB
+def load_csv_to_chroma(csv_path):
+    df = pd.read_csv(csv_path)
+
+    # Ensure the embedding column is properly formatted
+    df['embedding'] = df['embedding'].apply(lambda x: np.array(eval(x)).tolist())
+
+    # Create Chroma collection
+    collection_name = "video_subtitles"
+    if collection_name in [col.name for col in client.list_collections()]:
+        client.delete_collection(name=collection_name)
+
+    collection = client.create_collection(name=collection_name)
+
+    # Add data to Chroma DB
+    for i, row in df.iterrows():
+        collection.add(
+            ids=[str(i)],
+            documents=[row['text']],
+            embeddings=[row['embedding']]
+        )
+
+    return collection
+
+# Handle CSV upload and save permanently
+if csv_file:
+    st.sidebar.success("CSV uploaded successfully!")
+    
+    # Save CSV permanently
+    csv_path = save_csv_permanently(csv_file)
+    st.sidebar.success(f"CSV saved permanently at: {csv_path}")
+
+    # Load into Chroma DB
+    with st.spinner("Loading CSV into Chroma DB..."):
+        collection = load_csv_to_chroma(csv_path)
+    st.sidebar.success("CSV loaded into Chroma DB ✅")
+
+# Whisper model for transcription
 whisper_model = pipeline("automatic-speech-recognition", model="openai/whisper-base")
 
 # Function to extract audio from video
@@ -18,92 +70,81 @@ def extract_audio(video_file, audio_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
+# Function to transcribe audio
 def transcribe_audio(audio_file):
     result = whisper_model(audio_file)
     return result['text']
 
-# Function to generate embeddings
+# Generate embeddings for the transcription
 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 search subtitles in chunks
-def search_in_chunks(transcribed_text, chunk_files, top_k=10):
-    query_embedding = np.array(generate_embedding(transcribed_text)).reshape(1, -1)
-    results = []
-
-    for chunk_file in chunk_files:
-        chunk = pd.read_csv(chunk_file)
-        subtitle_embeddings = np.array([np.array(eval(e)) for e in chunk['embedding'].tolist()])
-        
-        # Calculate cosine similarity
-        similarities = cosine_similarity(query_embedding, subtitle_embeddings).flatten()
-
-        for idx, similarity in enumerate(similarities):
-            results.append({
-                "text": chunk.iloc[idx]['text'],
-                "cosine_similarity": similarity
-            })
-
-    # Sort results by similarity score
-    results = sorted(results, key=lambda x: x['cosine_similarity'], reverse=True)[:top_k]
-    return results
+# Search subtitles in Chroma DB
+def search_in_chroma(transcribed_text, collection, top_k=10):
+    query_embedding = np.array(generate_embedding(transcribed_text))
+
+    # Query Chroma DB
+    results = collection.query(
+        query_embeddings=[query_embedding.tolist()],
+        n_results=top_k
+    )
+
+    # Prepare results with cosine similarity
+    subtitles = []
+    for i, doc in enumerate(results['documents'][0]):
+        embedding = np.array(results['embeddings'][0][i])
+        similarity = cosine_similarity([query_embedding], [embedding])[0][0]
+
+        subtitles.append({
+            "text": doc,
+            "cosine_similarity": similarity
+        })
+
+    # Sort results by similarity
+    subtitles = sorted(subtitles, key=lambda x: x['cosine_similarity'], reverse=True)
+    return subtitles
 
 # Streamlit UI
-st.title("🎥 Video Subtitle Search with Hugging Face Whisper and Chunking")
+st.title("🎥 Video Subtitle Search with Chroma DB")
 
 # Upload video
 uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov", "mkv"])
 
-if uploaded_file:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
-        temp_video.write(uploaded_file.getbuffer())
-        video_path = temp_video.name
+if uploaded_file and csv_file:
+    # Create temporary directory
+    temp_dir = tempfile.mkdtemp()
+
+    # Save video temporarily
+    video_path = os.path.join(temp_dir, "temp_video.mp4")
+    with open(video_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
 
-    audio_path = "temp_audio.wav"
+    audio_path = os.path.join(temp_dir, "temp_audio.wav")
 
     # Extract audio
     st.info("Extracting audio...")
     extract_audio(video_path, audio_path)
 
-    # Transcribe audio using Hugging Face Whisper
+    # Transcribe audio
     st.info("Transcribing audio...")
     transcribed_text = transcribe_audio(audio_path)
     st.text_area("Transcribed Text", transcribed_text, height=150)
 
-    # 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)
-
-    # Search subtitles in chunks
-    st.info("Searching subtitles in chunks...")
-    matching_subtitles = search_in_chunks(transcribed_text, chunk_files)
+    # Search in Chroma DB
+    st.info("Searching subtitles in Chroma DB...")
+    matching_subtitles = search_in_chroma(transcribed_text, collection)
 
     # Display video
     st.video(video_path)
 
-    # Display matching subtitles with similarity scores
-    st.subheader("📜 Matching Subtitles (Chunking + Cosine Similarity)")
+    # Display matching subtitles
+    st.subheader("📜 Matching Subtitles with Cosine Similarity")
     for sub in matching_subtitles:
         st.write(f"**Subtitle:** {sub['text']}")
         st.write(f"**Cosine Similarity:** {sub['cosine_similarity']:.4f}")
         st.write("---")
 
-    # Cleanup
-    os.remove(video_path)
-    os.remove(audio_path)
-
-    # Remove chunk files
-    for chunk_file in chunk_files:
-        os.remove(chunk_file)
+    # Cleanup temporary files and directory
+    shutil.rmtree(temp_dir)
+    st.success("Temporary files cleaned up successfully ✅")