diff --git a/src/dataset.py b/src/dataset.py deleted file mode 100644 index b7c3b5f..0000000 --- a/src/dataset.py +++ /dev/null @@ -1,88 +0,0 @@ -import jsonlines -import numpy as np -from typing import List, Dict, Tuple -import openai - -from settings import DATA_PATH, EMBEDDING_MODEL -import config - -from helper import split_article, split_into_sentences - -openai.api_key = config.OPENAI_API_KEY - - -class Dataset: - def __init__(self, - path: str, # Path to the dataset .jsonl file. - sources: List[str] = None, # List of sources to include. If None, include all sources. - max_paragraph_length: Tuple[int, int] = None # (max number of words in a paragraph, max number of characters in a paragraph) (TBD) - ): - - self.path = path - self.sources = sources - self.max_paragraph_length = max_paragraph_length - - self.data: List[Tuple[str, str, str]] = [] # List of tuples, each containing the title of an article, its URL, and text. E.g.: [('title', 'url', 'text'), ...] - self.embed_split: List[str] = [] # List of strings, each being a few paragraphs from a single article (not exceeding 1000 words). - - self.num_articles: Dict[str, int] = {} # Number of articles per source. E.g.: {'source1': 10, 'source2': 20, 'total': 30} - if sources is None: - self.num_articles['total'] = 0 - else: - for source in sources: - self.num_articles[source] = 0 - self.num_articles['total'] = 0 - - self.total_char_count = 0 - self.total_word_count = 0 - self.total_sentence_count = 0 - self.total_paragraph_count = 0 - - def get_alignment_texts(self): - with jsonlines.open(self.path, "r") as reader: - for entry in reader: - try: - if self.sources is None: - if entry['source'] not in self.num_articles: - self.num_articles[entry['source']] = 1 - else: - self.num_articles[entry['source']] += 1 - self.num_articles['total'] += 1 - else: - if entry['source'] in self.sources: - self.num_articles[entry['source']] += 1 - self.num_articles['total'] += 1 - else: - continue - - # BIG PROBLEM: Very often, the post will have no URL, so this will fail. (TODO: Fix this.) - self.data.append((entry['title'], entry['url'], entry['text'])) - paragraphs = split_article(entry['text']) - self.embed_split.extend(paragraphs) - - self.total_char_count += len(entry['text']) - self.total_word_count += len(entry['text'].split()) - self.total_sentence_count += len(split_into_sentences(entry['text'])) - self.total_paragraph_count += len(paragraphs) - except KeyError: - pass - - def get_embedding(text: str) -> np.ndarray: - result = openai.Embedding.create(model=EMBEDDING_MODEL, input=text) - return result["data"][0]["embedding"] - - def get_embeddings(self): - self.embeddings = np.array([self.get_embedding(text) for text in self.embed_split]) - - def save_embeddings(self, path: str): - np.save(path, self.embeddings) - - def load_embeddings(self, path: str): - self.embeddings = np.load(path) - - -if __name__ == "__main__": - dataset = Dataset(DATA_PATH) - dataset.get_alignment_texts() - dataset.get_embeddings() - dataset.save_embeddings("embeddings.npy") diff --git a/src/helper.py b/src/helper.py deleted file mode 100644 index e1bab8d..0000000 --- a/src/helper.py +++ /dev/null @@ -1,61 +0,0 @@ -import re -from typing import List - -# FROM https://stackoverflow.com/a/31505798/16185542 -# -*- coding: utf-8 -*- -alphabets= "([A-Za-z])" -prefixes = "(Mr|St|Mrs|Ms|Dr)[.]" -suffixes = "(Inc|Ltd|Jr|Sr|Co)" -starters = "(Mr|Mrs|Ms|Dr|Prof|Capt|Cpt|Lt|He\s|She\s|It\s|They\s|Their\s|Our\s|We\s|But\s|However\s|That\s|This\s|Wherever)" -acronyms = "([A-Z][.][A-Z][.](?:[A-Z][.])?)" -websites = "[.](com|net|org|io|gov|edu|me)" -digits = "([0-9])" - -def split_into_sentences(text): - text = " " + text + " " - text = text.replace("\n"," ") - text = text.replace("?!", "?") - text = re.sub(prefixes,"\\1",text) - text = re.sub(websites,"\\1",text) - text = re.sub(digits + "[.]" + digits,"\\1\\2",text) - if "..." in text: text = text.replace("...","") - if "Ph.D" in text: text = text.replace("Ph.D.","PhD") - text = re.sub("\s" + alphabets + "[.] "," \\1 ",text) - text = re.sub(acronyms+" "+starters,"\\1 \\2",text) - text = re.sub(alphabets + "[.]" + alphabets + "[.]" + alphabets + "[.]","\\1\\2\\3",text) - text = re.sub(alphabets + "[.]" + alphabets + "[.]","\\1\\2",text) - text = re.sub(" "+suffixes+"[.] "+starters," \\1 \\2",text) - text = re.sub(" "+suffixes+"[.]"," \\1",text) - text = re.sub(" " + alphabets + "[.]"," \\1",text) - if "”" in text: text = text.replace(".”","”.") - if "\"" in text: text = text.replace(".\"","\".") - if "!" in text: text = text.replace("!\"","\"!") - if "?" in text: text = text.replace("?\"","\"?") - text = text.replace(".",".") - text = text.replace("?","?") - text = text.replace("!","!") - text = text.replace("",".") - - sentences = text.split("") - sentences = sentences[:-1] - sentences = [s.strip() for s in sentences] - - if sentences == []: - sentences = [text.strip()] - return sentences - - -def split_article(text: str) -> List[str]: # THIS IS COMPLETELY BROKEN AND WRONG. TODO: FIX IT. - # Receives one text (str) and returns a list of sections (List[str]), each section being a few appended paragraphs that do not exceed 1000 words. - # This is done to avoid the 8000 token limit of OpenAI embeddings. - sections = [] - section = "" - paragraphs = text.split('\n') - for paragraph in paragraphs: - if paragraph == "": continue - if len(section.split()) + len(paragraph.split()) > 1000 or len(section) + len(paragraph) > 7000: - sections.append(section) - section = "" - section += f"{paragraph}\n" - sections.append(section) - return sections \ No newline at end of file