Update app.py

app.py

@@ -9,8 +9,12 @@ from sentence_transformers import (
     SentenceTransformer
 )
 
+from huggingface_hub import (
+    InferenceClient
+)
+
 # =====================================================
-# BUILD DATABASE IF NEEDED
+# BUILD DB IF NEEDED
 # =====================================================
 
 def build_database_if_needed():
@@ -27,37 +31,24 @@ def build_database_if_needed():
         index_exists and metadata_exists
     ):
 
-        print(
-            "Building database..."
-        )
-
         subprocess.run(
-            ["python", "build_database.py"]
-        )
-
-        return (
-            "⚡ Building database..."
+            ["python",
+             "build_database.py"]
         )
 
-    return (
-        "✅ Database Ready"
-    )
-
-status_message = (
-    build_database_if_needed()
-)
+build_database_if_needed()
 
 # =====================================================
-# LOAD MODELS
+# LOAD EMBEDDING MODEL
 # =====================================================
 
-print("Loading embedding model...")
-
 model = SentenceTransformer(
     "all-MiniLM-L6-v2"
 )
 
-print("Loading vector DB...")
+# =====================================================
+# LOAD VECTOR DB
+# =====================================================
 
 index = faiss.read_index(
     "vector_store/faiss_index.bin"
@@ -71,16 +62,28 @@ with open(
     metadata = pickle.load(f)
 
 # =====================================================
-# SEARCH ENGINE
+# LOAD LLAMA 3
 # =====================================================
 
-def search_entertainment(query):
+hf_token = os.getenv(
+    "HF_TOKEN"
+)
 
-    if not query.strip():
+client = InferenceClient(
+    model=
+    "meta-llama/Meta-Llama-3-8B-Instruct",
 
-        return (
-            "Please ask a question."
-        )
+    token=hf_token
+)
+
+# =====================================================
+# AI SEARCH
+# =====================================================
+
+def entertainment_ai(query):
+
+    if not query.strip():
+        return "Please ask something."
 
     query_embedding = model.encode(
         [query]
@@ -95,99 +98,100 @@ def search_entertainment(query):
         5
     )
 
-    retrieved_docs = []
+    retrieved_context = []
 
     for idx in indices[0]:
 
         doc = metadata[idx]
 
-        retrieved_docs.append(
-            f"""
-🎬 Series:
-{doc['series']}
-
-📺 Season:
-{doc['season']}
-
-🎞 Episode:
-{doc['episode']}
-
-📝 Context:
-{doc['context']}
-"""
+        retrieved_context.append(
+            doc["context"]
         )
 
-    return "\n\n".join(
-        retrieved_docs
+    combined_context = (
+        "\n\n---\n\n".join(
+            retrieved_context
+        )
     )
 
-# =====================================================
-# 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;
+    messages = [
+        {
+            "role":
+            "system",
+
+            "content":
+            """
+            You are Universal
+            Entertainment AI.
+
+            Rules:
+            - Use ONLY retrieved context
+            - Never invent lore
+            - Be accurate
+            - Explain naturally
+            """
+        },
+
+        {
+            "role":
+            "user",
+
+            "content":
+            f"""
+            Question:
+            {query}
+
+            Context:
+            {combined_context}
+            """
+        }
+    ]
+
+    try:
+
+        response = (
+            client.chat.completions.create(
+                messages=messages,
+                max_tokens=250,
+                temperature=0.2
+            )
+        )
 
-    font-size:
-    42px;
+        return (
+            response
+            .choices[0]
+            .message.content
+        )
 
-    font-weight:
-    bold;
+    except Exception as e:
 
-    color:
-    cyan;
-}
-"""
+        return (
+            f"Error: {str(e)}"
+        )
 
-with gr.Blocks(
-    css=custom_css
-) as demo:
+# =====================================================
+# UI
+# =====================================================
 
-    gr.HTML("""
-    <div class="title">
-    🎬 Universal Entertainment AI
-    </div>
-    """)
+with gr.Blocks() as demo:
 
     gr.Markdown(
-        "Ask anything about TV shows."
+        "# 🎬 Universal Entertainment AI"
     )
 
-    gr.Textbox(
-        value=status_message,
-        label="System Status"
+    gr.Markdown(
+        "Ask anything about TV shows."
     )
 
     query = gr.Textbox(
-        label="Ask Question",
+        label="Question",
         placeholder=
-        "Example: Why did Homelander betray Starlight?"
+        "Why does Reddington want Elizabeth?"
     )
 
     output = gr.Textbox(
-        label="AI Response",
-        lines=20
+        label="AI Answer",
+        lines=12
     )
 
     button = gr.Button(
@@ -195,7 +199,7 @@ with gr.Blocks(
     )
 
     button.click(
-        fn=search_entertainment,
+        entertainment_ai,
         inputs=query,
         outputs=output
     )