Moved to Pinecode instead of pickle

.gitignore

@@ -131,14 +131,15 @@ dmypy.json
 # Other
 *alignment_texts.jsonl
 *config.py
-*.npy
 *.DS_Store
-src/tmp.py
+*.env
+*.npy
+*.pkl
 
 .vercel/
-api/dataset.pkl
 temp/
+*tmp.py
 
+api/dataset.pkl
 api/dataset_big.pkl
-
 api/dataset_300.pkl

api/chat.py

@@ -114,19 +114,20 @@ def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Bl
 
     return prompt
 
-# ------------------------------------------------------------------------------
+# ------------------------------- completion code -------------------------------
 
-# returns either (True, reply string, embeddings) or (False, error message string, None)
-def talk_to_robot(dataset_dict, query: str, history: List[Dict[str, str]], k: int = 10):
+# returns either (True, reply string, top_k_blocks)) or (False, error message string, None)
+def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = 10):
 
 
     # 1. Find the most relevant blocks from the Alignment Research Dataset
-    top_k_blocks: List[Block] = get_top_k_blocks(dataset_dict, query, k)
+    top_k_blocks = get_top_k_blocks(index, query, k)
 
 
     # 2. Generate a prompt
     prompt = construct_prompt(query, history, top_k_blocks)
 
+    
     if DEBUG_PRINT:
         print('\n' * 10)
         print(" ------------------------------ prompt: -----------------------------")

api/get_blocks.py

@@ -44,8 +44,9 @@ def get_embedding(text: str) -> np.ndarray:
 
             time.sleep(min(max_wait_time, 2 ** attempt))
 
-# Get the k blocks most semantically similar to the query.
-def get_top_k_blocks(data, user_query: str, k: int = 10) -> List[Block]:
+
+# Get the k blocks most semantically similar to the query using Pinecone.
+def get_top_k_blocks(index, user_query: str, k: int = 10) -> List[Block]:
 
     # print time
     t = time.time()
@@ -56,24 +57,27 @@ def get_top_k_blocks(data, user_query: str, k: int = 10) -> List[Block]:
     t1 = time.time()
     print("Time to get embedding: ", t1 - t)
 
-    similarity_scores = np.dot(data["embeddings"], query_embedding) # big fat calculation
-
+    query_response = index.query(
+        namespace="alignment-search",  # ugly, sorry
+        top_k=k,
+        include_values=False,
+        include_metadata=True,
+        vector=query_embedding
+    )
+    blocks = [
+        Block(
+            title = match['metadata']['title'],
+            author = match['metadata']['author'],
+            date = match['metadata']['date'],
+            url = match['metadata']['url'],
+            tags = match['metadata']['tags'],
+            text = match['metadata']['text']
+        ) for match in query_response['matches']
+    ]
     t2 = time.time()
-    print("Time to get similarity scores: ", t2 - t1)
+
+    print("Time to get top-k blocks: ", t2 - t1)
     
-    top_k_block_indices = list(reversed(np.argpartition(similarity_scores, -k)[-k:])) # Get the top k indices of the blocks
-
-    t3 = time.time()
-    print("Time to get top k indices: ", t3 - t2)
-
-    top_k_metadata_indexes = [data["embeddings_metadata_index"][i] for i in top_k_block_indices]
-    top_k_texts = [strip_block(data["embedding_strings"][i]) for i in top_k_block_indices]
-    top_k_metadata = [data["metadata"][i] for i in top_k_metadata_indexes]
-
-    # Combine the top k texts and metadata into a list of Block objects
-    top_k_metadata_and_text = [list(top_k_metadata[i]) + [top_k_texts[i]] for i in range(len(top_k_metadata))]
-    blocks = [Block(*block) for block in top_k_metadata_and_text]
-
     # for all blocks that are "the same" (same title, author, date, url, tags),
     # combine their text with "....." in between. Return them in order such
     # that the combined block has the minimum index of the blocks combined.

api/main.py

@@ -5,10 +5,12 @@ from chat import talk_to_robot
 import dataclasses
 import os
 import openai
-import pickle
+import pinecone
+
 
 # ---------------------------------- env setup ---------------------------------
 
+
 if os.path.exists('.env'):
     from dotenv import load_dotenv
     load_dotenv()
@@ -16,13 +18,14 @@ if os.path.exists('.env'):
 OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
 openai.api_key = OPENAI_API_KEY
 
-# -------------------------------- load dataset --------------------------------
-
-
-print('Loading dataset...')
-with open('dataset.pkl', 'rb') as f:
-    dataset_dict = pickle.load(f)
-print('Done!')
+PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
+PINECONE_ENV = "us-east1-gcp"
+pinecone.init(
+    api_key=PINECONE_API_KEY,
+    environment=PINECONE_ENV
+)
+INDEX_NAME = "alignment-search"
+index = pinecone.Index(index_name=INDEX_NAME)
 
 
 # ---------------------------------- web setup ---------------------------------
@@ -31,7 +34,6 @@ app = Flask(__name__)
 cors = CORS(app)
 app.config['CORS_HEADERS'] = 'Content-Type'
 
-
 # ------------------------------- semantic search ------------------------------
 
 
@@ -39,7 +41,7 @@ app.config['CORS_HEADERS'] = 'Content-Type'
 @cross_origin()
 def semantic():
     query = request.json['query']
-    return jsonify([dataclasses.asdict(block) for block in get_top_k_blocks(dataset_dict, query)])
+    return jsonify([dataclasses.asdict(block) for block in get_top_k_blocks(index, query)])
 
 
 # ------------------------------------ chat ------------------------------------
@@ -52,13 +54,16 @@ def chat():
     query = request.json['query']
     history = request.json['history']
 
-    is_valid, response, context = talk_to_robot(dataset_dict, query, history)
+    is_valid, response, context = talk_to_robot(index, query, history)
 
     if is_valid:
         return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
     else:
         return jsonify({'error': response})
 
+
 # ------------------------------------------------------------------------------
 
-if __name__ == '__main__': app.run(debug=True, port=3000)
+
+if __name__ == '__main__':
+    app.run(debug=True, port=3000)

api/requirements.txt

@@ -17,3 +17,4 @@ openai==0.27.2
 numpy==1.24.2
 tenacity==8.2.2
 tiktoken
+pinecone-client