CoolFace
Apppublic

Atulkumar001/Universal-Entertainment-AI

sourceHugging Facemitupdated 4mo agoView on Hugging Face
0likes
build_database.py329 linesDownload Raw Back to root
1import os2import pickle3import numpy as np4import pandas as pd5import faiss6 7from sentence_transformers import (8    SentenceTransformer9)10 11# =====================================================12# CREATE FOLDERS13# =====================================================14 15os.makedirs(16    "vector_store",17    exist_ok=True18)19 20# =====================================================21# TV KNOWLEDGE DATABASE22# =====================================================23 24documents = [25 26# =====================================================27# THE BLACKLIST28# =====================================================29 30{31"series":"The Blacklist",32"season":1,33"episode":1,34"title":"Pilot",35"context":36"""37Reddington surrenders to the FBI38and demands to work only with39Elizabeth Keen.40"""41},42 43{44"series":"The Blacklist",45"season":1,46"episode":1,47"title":"Pilot",48"context":49"""50Reddington introduces51the blacklist,52a secret list of dangerous53criminals.54"""55},56 57{58"series":"The Blacklist",59"season":1,60"episode":1,61"title":"Pilot",62"context":63"""64Ranko Zamani is identified65as a dangerous criminal66targeting a general's daughter.67"""68},69 70# =====================================================71# BREAKING BAD72# =====================================================73 74{75"series":"Breaking Bad",76"season":1,77"episode":1,78"title":"Pilot",79"context":80"""81Walter White discovers he has82lung cancer and starts making83methamphetamine to support84his family.85"""86},87 88{89"series":"Breaking Bad",90"season":2,91"episode":8,92"title":"Better Call Saul",93"context":94"""95Saul Goodman becomes Walter96White's lawyer and helps him97expand his illegal drug business.98"""99},100 101{102"series":"Breaking Bad",103"season":5,104"episode":16,105"title":"Felina",106"context":107"""108Walter White sacrifices himself109after rescuing Jesse Pinkman.110"""111},112 113# =====================================================114# THE BOYS115# =====================================================116 117{118"series":"The Boys",119"season":1,120"episode":1,121"title":"The Name of the Game",122"context":123"""124Starlight joins The Seven while125Homelander is introduced as126its dangerous leader.127"""128},129 130{131"series":"The Boys",132"season":2,133"episode":8,134"title":"What I Know",135"context":136"""137Starlight helps expose138Homelander and Vought.139"""140},141 142{143"series":"The Boys",144"season":3,145"episode":8,146"title":"The Instant White-Hot Wild",147"context":148"""149Butcher and Homelander150continue their rivalry.151"""152},153 154# =====================================================155# GAME OF THRONES156# =====================================================157 158{159"series":"Game of Thrones",160"season":1,161"episode":9,162"title":"Baelor",163"context":164"""165Ned Stark is publicly executed166on King Joffrey's orders.167"""168},169 170{171"series":"Game of Thrones",172"season":3,173"episode":9,174"title":"The Rains of Castamere",175"context":176"""177The Red Wedding results in178the deaths of Robb Stark,179Talisa and Catelyn Stark.180"""181},182 183{184"series":"Game of Thrones",185"season":8,186"episode":6,187"title":"The Iron Throne",188"context":189"""190Jon Snow kills Daenerys191to stop her violent rule.192"""193},194 195# =====================================================196# FRIENDS197# =====================================================198 199{200"series":"Friends",201"season":1,202"episode":1,203"title":"Pilot",204"context":205"""206Rachel leaves Barry at the altar207and reconnects with Monica.208"""209},210 211{212"series":"Friends",213"season":5,214"episode":14,215"title":"Everybody Finds Out",216"context":217"""218The group discovers Monica and219Chandler are secretly dating.220"""221},222 223{224"series":"Friends",225"season":10,226"episode":18,227"title":"The Last One",228"context":229"""230Ross and Rachel reunite.231The friends leave the apartment.232"""233},234 235# =====================================================236# STRANGER THINGS237# =====================================================238 239{240"series":"Stranger Things",241"season":1,242"episode":1,243"title":"Chapter One",244"context":245"""246Will Byers disappears and Eleven247appears mysteriously.248"""249},250 251{252"series":"Stranger Things",253"season":4,254"episode":9,255"title":"The Piggyback",256"context":257"""258Eleven battles Vecna while259Max is critically injured.260"""261}262]263 264# =====================================================265# DATAFRAME266# =====================================================267 268df = pd.DataFrame(269    documents270)271 272print(273    f"Loaded {len(df)} entries"274)275 276# =====================================================277# EMBEDDINGS278# =====================================================279 280model = SentenceTransformer(281    "all-MiniLM-L6-v2"282)283 284texts = df["context"].tolist()285 286embeddings = model.encode(287    texts288)289 290embeddings = np.array(291    embeddings292).astype("float32")293 294# =====================================================295# FAISS296# =====================================================297 298dimension = embeddings.shape[1]299 300index = faiss.IndexFlatL2(301    dimension302)303 304index.add(305    embeddings306)307 308# =====================================================309# SAVE310# =====================================================311 312faiss.write_index(313    index,314    "vector_store/faiss_index.bin"315)316 317with open(318    "vector_store/metadata.pkl",319    "wb"320) as f:321 322    pickle.dump(323        documents,324        f325    )326 327print(328    "โœ… Entertainment database built"329)