From fb34d1b95c51615829220fba30248e3a14958bb1 Mon Sep 17 00:00:00 2001 From: henri123lemoine Date: Sat, 25 Mar 2023 15:47:39 -0400 Subject: [PATCH] Modified embeddings.py's outputs to be more informative. --- src/dataset.py | 4 +-- web/api/embeddings.py | 63 ++++++++++++++++++++++++------------------- 2 files changed, 38 insertions(+), 29 deletions(-) diff --git a/src/dataset.py b/src/dataset.py index 5468717..acfdbe3 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -59,8 +59,8 @@ class Dataset: self.min_tokens_per_block = min_tokens_per_block # for the text splitter self.max_tokens_per_block = max_tokens_per_block # for the text splitter - self.metadata: List[Tuple[str]] = [] # List of tuples, each containing the title of an article, its URL, and text. E.g.: [('title', 'url', 'text'), ...] - self.embedding_strings: List[str] = [] # List of strings, each being a few paragraphs from a single article (not exceeding 1000 words). + self.metadata: List[Tuple[str]] = [] # List of tuples, each containing the title, author, date, URL, and tags of an article. + self.embedding_strings: List[str] = [] # List of strings, each being a few paragraphs from a single article (not exceeding max_tokens_per_block tokens). self.embeddings_metadata_index: List[int] = [] # List of integers, each being the index of the article from which the embedding string was taken. self.articles_count: DefaultDict[str, int] = defaultdict(int) # Number of articles per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} diff --git a/web/api/embeddings.py b/web/api/embeddings.py index f071216..07f0168 100644 --- a/web/api/embeddings.py +++ b/web/api/embeddings.py @@ -15,7 +15,7 @@ class handler(BaseHTTPRequestHandler): post_data = self.rfile.read(content_length) data = json.loads(post_data) - results = {}; + results = {} for i, link in enumerate(embeddings(data['query'])): results[i] = json.dumps(link.__dict__) @@ -23,12 +23,6 @@ class handler(BaseHTTPRequestHandler): self.wfile.write(json.dumps(results).encode('utf-8')) -class Link: - def __init__(self, url, title): - self.url = url - self.title = title - - # -------------------------------- non-web-code -------------------------------- import time import pickle @@ -65,6 +59,18 @@ PATH_TO_EMBEDDINGS = project_path / "web" / "api" / "data" / "embeddings.npy" # PATH_TO_DATASET = project_path / "web" / "api" / "data" / "dataset.pkl" # Path to the saved dataset (.pkl) file, containing the dataset class object. +class Dataset: + pass + +class Block: + def __init__(self, title: str, author: str, date: str, url: str, tags: str, text: str): + self.title = title + self.author = author + self.date = date + self.url = url + self.tags = tags + self.text = text + def retry_on_exception_types(exception_types: List[Type[Exception]], stop_after_attempt: int, max_wait_time: int) -> Callable: def decorator(func: Callable) -> Callable: @wraps(func) @@ -98,10 +104,7 @@ def get_embedding(text: str) -> np.ndarray: ) return result["data"][0]["embedding"] -class Dataset: - pass - -def get_top_k_blocks(user_query: str, k: int, HyDE: bool = False): +def get_top_k_blocks(user_query: str, k: int, HyDE: bool = False) -> List[Block]: """Get the top k blocks that are most semantically similar to the query, using the provided dataset. Args: @@ -110,10 +113,9 @@ def get_top_k_blocks(user_query: str, k: int, HyDE: bool = False): HyDE (bool, optional): Whether to use HyDE or not. Defaults to False. Returns: - List[str]: A list of the top k blocks that are most semantically similar to the query. + List[Block]: A list of the top k blocks that are most semantically similar to the query. """ # Get the dataset - print(f"\n\nLoading {PATH_TO_DATASET}...\n\n") with open(PATH_TO_DATASET, "rb") as f: metadataset = pickle.load(f) @@ -139,30 +141,37 @@ def get_top_k_blocks(user_query: str, k: int, HyDE: bool = False): similarity_scores = np.dot(metadataset.embeddings, query_embedding) ordered_blocks = np.argsort(similarity_scores)[::-1] # Sort the blocks by similarity score - top_k_indices = ordered_blocks[:k] # Get the top k indices + top_k_text_indices = ordered_blocks[:k] # Get the top k indices of the blocks + top_k_metadata_indexes = [metadataset.embeddings_metadata_index[i] for i in top_k_text_indices] - # Get associated strings - top_k_strings = [metadataset.embedding_strings[i] for i in top_k_indices] # Get the top k strings + # Get the top k blocks (title, author, date, url, tags, text) + top_k_texts = [metadataset.embedding_strings[i] for i in top_k_text_indices] # Get the top k texts + top_k_metadata = [metadataset.metadata[i] for i in top_k_metadata_indexes] # Get the top k metadata (title, author, date, url, tags) - # Get associated links - top_k_links = [metadataset.metadata[metadataset.embeddings_metadata_index[i]][3] for i in top_k_indices] # Get the top k sources + # 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(k)] - links = [] - for string, link in zip(top_k_strings, top_k_links): - links.append(Link(link, string)) + top_k_blocks = [Block(*block) for block in top_k_metadata_and_text] - return links + return top_k_blocks def embeddings(query): # write a function here that takes a query, returns a bunch of semantically similar links - link_list = get_top_k_blocks(query, 8, HyDE=False) + top_k_blocks = get_top_k_blocks(query, 8, HyDE=False) - return link_list + return top_k_blocks if __name__ == "__main__": # Test the embeddings function - links = embeddings("The last enemy that shall be destroyed is death.") - for link in links: - print(link.title, link.url) \ No newline at end of file + blocks = embeddings("Artificial Intelligence stinks.") + for link in blocks: + print(f"Title: {link.title}") + print(f"Author: {link.author}") + print(f"Date: {link.date}") + print(f"URL: {link.url}") + print(f"Tags: {link.tags}") + print(f"Text: {link.text}") + print() + print() \ No newline at end of file