From 7bd9c1e3a4ec9997fd4bbda5e595ab8df7024e1a Mon Sep 17 00:00:00 2001 From: navanchauhan Date: Sat, 21 May 2022 10:54:12 -0600 Subject: [PATCH] updated create and update embeddings --- db2pc.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/db2pc.py b/db2pc.py index f0d6ba9..5f28167 100644 --- a/db2pc.py +++ b/db2pc.py @@ -1,16 +1,37 @@ +import os +import pinecone + from database import * import pandas as pd from sentence_transformers import SentenceTransformer -database_url = "sqlite:///jlm.db" +from tqdm import tqdm +database_url = "sqlite:///jlm.db" engine, Session = init_db_stuff(database_url) +PINECONE_KEY = os.getenv("PINECONE_API_DEFAULT") +pinecone.init(api_key=PINECONE_KEY, environment="us-west1-gcp") +index = pinecone.Index("movies") + model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2") +batch_size = 32 + df = pd.read_sql("Select * from movies", engine) df["combined_text"] = df["title"] + ": " + df["overview"].fillna('') + " - " + df["tagline"].fillna('') + " Genres:- " + df["genres"].fillna('') print(len(df["combined_text"].tolist())) +for x in tqdm(range(0,len(df),batch_size)): + to_send = [] + trakt_ids = df["trakt_id"][x:x+batch_size].tolist() + sentences = df["combined_text"][x:x+batch_size].tolist() + embeddings = model.encode(sentences) + for idx, value in enumerate(trakt_ids): + to_send.append( + { + value: embeddings[idx] + }) + index.upsert(to_send)