Update build_database.py

build_database.py

@@ -1 +1,310 @@
-print("Database builder ready")
+import os
+import re
+import pickle
+import requests
+import pandas as pd
+import numpy as np
+import faiss
+
+from bs4 import BeautifulSoup
+from sentence_transformers import SentenceTransformer
+
+# =====================================================
+# CREATE FOLDERS
+# =====================================================
+
+os.makedirs(
+    "vector_store",
+    exist_ok=True
+)
+
+os.makedirs(
+    "data/processed",
+    exist_ok=True
+)
+
+# =====================================================
+# TV SHOW DATABASE
+# =====================================================
+
+tv_shows = {
+
+    "The Blacklist": [
+        {
+            "season": 1,
+            "episode": 1,
+            "title": "Pilot",
+            "url":
+            "https://subslikescript.com/series/The_Blacklist-2741602/season-1/episode-1-Pilot"
+        }
+    ],
+
+    "Breaking Bad": [
+        {
+            "season": 1,
+            "episode": 1,
+            "title": "Pilot",
+            "url":
+            "https://subslikescript.com/series/Breaking_Bad-903747/season-1/episode-1-Pilot"
+        }
+    ],
+
+    "The Boys": [
+        {
+            "season": 1,
+            "episode": 1,
+            "title":
+            "The Name of the Game",
+
+            "url":
+            "https://subslikescript.com/series/The_Boys-1190634/season-1/episode-1-The_Name_of_the_Game"
+        }
+    ]
+}
+
+# =====================================================
+# SCRAPE SCRIPT
+# =====================================================
+
+def scrape_script(url):
+
+    try:
+
+        response = requests.get(
+            url,
+            headers={
+                "User-Agent":
+                "Mozilla/5.0"
+            },
+            timeout=15
+        )
+
+        soup = BeautifulSoup(
+            response.text,
+            "html.parser"
+        )
+
+        script_div = soup.find(
+            "div",
+            class_="full-script"
+        )
+
+        if script_div:
+
+            return script_div.get_text(
+                separator="\n"
+            )
+
+        return None
+
+    except Exception as e:
+
+        print(
+            f"Error scraping {url}"
+        )
+
+        print(e)
+
+        return None
+
+
+# =====================================================
+# CLEAN SCRIPT
+# =====================================================
+
+def clean_script(raw_text):
+
+    lines = raw_text.split("\n")
+
+    cleaned_lines = []
+
+    dialogue_pattern = re.compile(
+        r"^([A-Z\s]+):\s*(.*)"
+    )
+
+    for line in lines:
+
+        line = line.strip()
+
+        if not line:
+            continue
+
+        match = dialogue_pattern.match(
+            line
+        )
+
+        if match:
+
+            character = (
+                match.group(1)
+                .strip()
+            )
+
+            dialogue = (
+                match.group(2)
+                .strip()
+            )
+
+            cleaned_lines.append(
+                f"{character}: {dialogue}"
+            )
+
+    return cleaned_lines
+
+
+# =====================================================
+# CREATE CHUNKS
+# =====================================================
+
+def create_chunks(lines):
+
+    window_size = 2
+
+    chunks = []
+
+    for i in range(len(lines)):
+
+        start = max(
+            0,
+            i - window_size
+        )
+
+        end = min(
+            len(lines),
+            i + window_size + 1
+        )
+
+        chunk = "\n".join(
+            lines[start:end]
+        )
+
+        chunks.append(chunk)
+
+    return chunks
+
+
+# =====================================================
+# BUILD DATASET
+# =====================================================
+
+all_documents = []
+
+for show_name, episodes in tv_shows.items():
+
+    print(
+        f"\nScraping {show_name}"
+    )
+
+    for ep in episodes:
+
+        raw_text = scrape_script(
+            ep["url"]
+        )
+
+        if not raw_text:
+            continue
+
+        lines = clean_script(
+            raw_text
+        )
+
+        chunks = create_chunks(
+            lines
+        )
+
+        for chunk in chunks:
+
+            all_documents.append({
+
+                "series":
+                show_name,
+
+                "season":
+                ep["season"],
+
+                "episode":
+                ep["episode"],
+
+                "title":
+                ep["title"],
+
+                "context":
+                chunk
+            })
+
+# =====================================================
+# DATAFRAME
+# =====================================================
+
+df = pd.DataFrame(
+    all_documents
+)
+
+print(
+    f"\nTotal chunks: {len(df)}"
+)
+
+df.to_csv(
+    "data/processed/shows.csv",
+    index=False
+)
+
+# =====================================================
+# EMBEDDINGS
+# =====================================================
+
+print(
+    "\nLoading embedding model..."
+)
+
+model = SentenceTransformer(
+    "all-MiniLM-L6-v2"
+)
+
+texts = df["context"].tolist()
+
+embeddings = model.encode(
+    texts,
+    show_progress_bar=True
+)
+
+embeddings = np.array(
+    embeddings
+).astype("float32")
+
+# =====================================================
+# FAISS INDEX
+# =====================================================
+
+dimension = embeddings.shape[1]
+
+index = faiss.IndexFlatL2(
+    dimension
+)
+
+index.add(
+    embeddings
+)
+
+# =====================================================
+# SAVE VECTOR DB
+# =====================================================
+
+faiss.write_index(
+    index,
+    "vector_store/faiss_index.bin"
+)
+
+with open(
+    "vector_store/metadata.pkl",
+    "wb"
+) as f:
+
+    pickle.dump(
+        all_documents,
+        f
+    )
+
+print(
+    "\n✅ Entertainment database built successfully!"
+)