diff --git a/.gitignore b/.gitignore index df62ec6..80e3f6c 100644 --- a/.gitignore +++ b/.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 diff --git a/api/chat.py b/api/chat.py index 09e028d..9efb2ab 100644 --- a/api/chat.py +++ b/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: -----------------------------") diff --git a/api/get_blocks.py b/api/get_blocks.py index 5f6b3dd..7368bcd 100644 --- a/api/get_blocks.py +++ b/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. diff --git a/api/main.py b/api/main.py index 662b507..099c2ca 100644 --- a/api/main.py +++ b/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) \ No newline at end of file diff --git a/api/requirements.txt b/api/requirements.txt index a0f4e56..b819d83 100644 --- a/api/requirements.txt +++ b/api/requirements.txt @@ -17,3 +17,4 @@ openai==0.27.2 numpy==1.24.2 tenacity==8.2.2 tiktoken +pinecone-client