Update app.py

app.py

@@ -1,7 +1,17 @@
 import gradio as gr
 import os
 import subprocess
+import pickle
+import faiss
+import numpy as np
 
+from sentence_transformers import (
+    SentenceTransformer
+)
+
+# =====================================================
+# BUILD DATABASE IF NEEDED
+# =====================================================
 
 def build_database_if_needed():
 
@@ -26,31 +36,168 @@ def build_database_if_needed():
         )
 
         return (
-            "⚡ Building Entertainment Database..."
+            "⚡ Building database..."
         )
 
     return (
-        "✅ Entertainment Database Ready"
+        "✅ Database Ready"
     )
 
-
 status_message = (
     build_database_if_needed()
 )
 
-with gr.Blocks() as demo:
+# =====================================================
+# LOAD MODELS
+# =====================================================
 
-    gr.Markdown(
-        "# 🎬 Universal Entertainment AI"
+print("Loading embedding model...")
+
+model = SentenceTransformer(
+    "all-MiniLM-L6-v2"
+)
+
+print("Loading vector DB...")
+
+index = faiss.read_index(
+    "vector_store/faiss_index.bin"
+)
+
+with open(
+    "vector_store/metadata.pkl",
+    "rb"
+) as f:
+
+    metadata = pickle.load(f)
+
+# =====================================================
+# SEARCH ENGINE
+# =====================================================
+
+def search_entertainment(query):
+
+    if not query.strip():
+
+        return (
+            "Please ask a question."
+        )
+
+    query_embedding = model.encode(
+        [query]
+    )
+
+    query_embedding = np.array(
+        query_embedding
+    ).astype("float32")
+
+    distances, indices = index.search(
+        query_embedding,
+        5
+    )
+
+    retrieved_docs = []
+
+    for idx in indices[0]:
+
+        doc = metadata[idx]
+
+        retrieved_docs.append(
+            f"""
+🎬 Series:
+{doc['series']}
+
+📺 Season:
+{doc['season']}
+
+🎞 Episode:
+{doc['episode']}
+
+📝 Context:
+{doc['context']}
+"""
+        )
+
+    return "\n\n".join(
+        retrieved_docs
     )
 
+# =====================================================
+# FUTURISTIC UI
+# =====================================================
+
+custom_css = """
+body {
+    background:
+    #050816 !important;
+}
+
+.gradio-container {
+
+    background:
+    radial-gradient(
+    circle at top left,
+    #101935 0%,
+    #050816 60%,
+    #02030a 100%)
+    !important;
+
+    color:
+    white !important;
+}
+
+.title {
+    text-align:
+    center;
+
+    font-size:
+    42px;
+
+    font-weight:
+    bold;
+
+    color:
+    cyan;
+}
+"""
+
+with gr.Blocks(
+    css=custom_css
+) as demo:
+
+    gr.HTML("""
+    <div class="title">
+    🎬 Universal Entertainment AI
+    </div>
+    """)
+
     gr.Markdown(
-        "Recruiter-level TV & Movie Intelligence Engine"
+        "Ask anything about TV shows."
     )
 
-    status = gr.Textbox(
+    gr.Textbox(
         value=status_message,
         label="System Status"
     )
 
+    query = gr.Textbox(
+        label="Ask Question",
+        placeholder=
+        "Example: Why did Homelander betray Starlight?"
+    )
+
+    output = gr.Textbox(
+        label="AI Response",
+        lines=20
+    )
+
+    button = gr.Button(
+        "🚀 Analyze"
+    )
+
+    button.click(
+        fn=search_entertainment,
+        inputs=query,
+        outputs=output
+    )
+
 demo.launch()