toolkit/utils.py

"""
The widgets defines utility functions for loading data, text cleaning,
and indexing documents, as well as classes for handling document queries
and formatting chat history.
"""
import re
import pickle
import string
import logging
import configparser
from enum import Enum
from typing import List, Tuple, Union
import nltk
from nltk.stem import WordNetLemmatizer
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import torch
import tiktoken
from langchain.vectorstores import Chroma

from langchain.schema import Document, BaseMessage
from langchain.embeddings import HuggingFaceEmbeddings, HuggingFaceInstructEmbeddings
from langchain.embeddings.openai import OpenAIEmbeddings


tokenizer_name = tiktoken.encoding_for_model("gpt-3.5-turbo")
tokenizer = tiktoken.get_encoding(tokenizer_name.name)

# if nltk stopwords, punkt and wordnet are not downloaded, download it
try:
    nltk.data.find("corpora/stopwords")
except LookupError:
    nltk.download("stopwords")
try:
    nltk.data.find("tokenizers/punkt")
except LookupError:
    nltk.download("punkt")
try:
    nltk.data.find("corpora/wordnet")
except LookupError:
    nltk.download("wordnet")

ChatTurnType = Union[Tuple[str, str], BaseMessage]
_ROLE_MAP = {"human": "Human: ", "ai": "Assistant: "}


class Config:
    """Initializes configs."""

    def __init__(self, config_file):
        self.config = configparser.ConfigParser(interpolation=None)
        self.config.read(config_file)

        # Tokens
        self.openai_api_key = self.config.get("tokens", "OPENAI_API_KEY")
        self.anthropic_api_key = self.config.get("tokens", "ANTHROPIC_API_KEY")
        self.together_api_key = self.config.get("tokens", "TOGETHER_API_KEY")
        self.huggingface_token = self.config.get("tokens", "HUGGINGFACE_TOKEN")
        self.version = self.config.get("tokens", "VERSION")

        # Directory
        self.docs_dir = self.config.get("directory", "DOCS_DIR")
        self.db_dir = self.config.get("directory", "db_DIR")
        self.local_model_dir = self.config.get("directory", "LOCAL_MODEL_DIR")

        # Parameters
        self.model_name = self.config.get("parameters", "MODEL_NAME")
        self.temperature = self.config.getfloat("parameters", "TEMPURATURE")
        self.max_chat_history = self.config.getint("parameters", "MAX_CHAT_HISTORY")
        self.max_llm_context = self.config.getint("parameters", "MAX_LLM_CONTEXT")
        self.max_llm_generation = self.config.getint("parameters", "MAX_LLM_GENERATION")
        self.embedding_name = self.config.get("parameters", "EMBEDDING_NAME")

        self.n_gpu_layers = self.config.getint("parameters", "N_GPU_LAYERS")
        self.n_batch = self.config.getint("parameters", "N_BATCH")

        self.base_chunk_size = self.config.getint("parameters", "BASE_CHUNK_SIZE")
        self.chunk_overlap = self.config.getint("parameters", "CHUNK_OVERLAP")
        self.chunk_scale = self.config.getint("parameters", "CHUNK_SCALE")
        self.window_steps = self.config.getint("parameters", "WINDOW_STEPS")
        self.window_scale = self.config.getint("parameters", "WINDOW_SCALE")

        self.retriever_weights = [
            float(x.strip())
            for x in self.config.get("parameters", "RETRIEVER_WEIGHTS").split(",")
        ]
        self.first_retrieval_k = self.config.getint("parameters", "FIRST_RETRIEVAL_K")
        self.second_retrieval_k = self.config.getint("parameters", "SECOND_RETRIEVAL_K")
        self.num_windows = self.config.getint("parameters", "NUM_WINDOWS")

        # Logging
        self.logging_enabled = self.config.getboolean("logging", "enabled")
        self.logging_level = self.config.get("logging", "level")
        self.logging_filename = self.config.get("logging", "filename")
        self.logging_format = self.config.get("logging", "format")

        self.configure_logging()

    def configure_logging(self):
        """
        Configure the logger for each .py files.
        """

        if not self.logging_enabled:
            logging.disable(logging.CRITICAL + 1)
            return

        log_level = self.config.get("logging", "level")
        log_filename = self.config.get("logging", "filename")
        log_format = self.config.get("logging", "format")

        logging.basicConfig(level=log_level, filename=log_filename, format=log_format)


def configure_logger():
    """
    Configure the logger for each .py files.
    """
    config = configparser.ConfigParser(interpolation=None)
    config.read("configparser.ini")

    enabled = config.getboolean("logging", "enabled")

    if not enabled:
        logging.disable(logging.CRITICAL + 1)
        return

    log_level = config.get("logging", "level")
    log_filename = config.get("logging", "filename")
    log_format = config.get("logging", "format")

    logging.basicConfig(level=log_level, filename=log_filename, format=log_format)


def tiktoken_len(text):
    """token length function"""
    tokens = tokenizer.encode(text, disallowed_special=())
    return len(tokens)


def check_device():
    """Check if cuda or MPS is available, else fallback to CPU"""
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"
    return device


def choose_embeddings(embedding_name):
    """Choose embeddings for a given model's name"""
    try:
        if embedding_name == "openAIEmbeddings":
            return OpenAIEmbeddings()
        elif embedding_name == "hkunlpInstructorLarge":
            device = check_device()
            return HuggingFaceInstructEmbeddings(
                model_name="hkunlp/instructor-large", model_kwargs={"device": device}
            )
        else:
            device = check_device()
            return HuggingFaceEmbeddings(model_name=embedding_name, device=device)
    except Exception as error:
        raise ValueError(f"Embedding {embedding_name} not supported") from error


def load_embedding(store_name, embedding, suffix, path):
    """Load chroma embeddings"""
    vector_store = Chroma(
        persist_directory=f"{path}/chroma_{store_name}_{suffix}",
        embedding_function=embedding,
    )
    return vector_store


def load_pickle(prefix, suffix, path):
    """Load langchain documents from a pickle file.

    Args:
        store_name (str): The name of the store where data is saved.
        suffix (str): Suffix to append to the store name.
        path (str): The path where the pickle file is stored.

    Returns:
        Document: documents from the pickle file
    """
    with open(f"{path}/{prefix}_{suffix}.pkl", "rb") as file:
        return pickle.load(file)


def clean_text(text):
    """
    Converts text to lowercase, removes punctuation, stopwords, and lemmatizes it
    for BM25 retriever.

    Parameters:
        text (str): The text to be cleaned.

    Returns:
        str: The cleaned and lemmatized text.
    """
    # remove [SEP] in the text
    text = text.replace("[SEP]", "")
    # Tokenization
    tokens = word_tokenize(text)
    # Lowercasing
    tokens = [w.lower() for w in tokens]
    # Remove punctuation
    table = str.maketrans("", "", string.punctuation)
    stripped = [w.translate(table) for w in tokens]
    # Keep tokens that are alphabetic, numeric, or contain both.
    words = [
        word
        for word in stripped
        if word.isalpha()
        or word.isdigit()
        or (re.search("\d", word) and re.search("[a-zA-Z]", word))
    ]
    # Remove stopwords
    stop_words = set(stopwords.words("english"))
    words = [w for w in words if w not in stop_words]
    # Lemmatization (or you could use stemming instead)
    lemmatizer = WordNetLemmatizer()
    lemmatized = [lemmatizer.lemmatize(w) for w in words]
    # Convert list of words to a string
    lemmatized_ = " ".join(lemmatized)

    return lemmatized_


class IndexerOperator(Enum):
    """
    Enumeration for different query operators used in indexing.
    """

    EQ = "=="
    GT = ">"
    GTE = ">="
    LT = "<"
    LTE = "<="


class DocIndexer:
    """
    A class to handle indexing and searching of documents.

    Attributes:
        documents (List[Document]): List of documents to be indexed.
    """

    def __init__(self, documents):
        self.documents = documents
        self.index = self.build_index(documents)

    def build_index(self, documents):
        """
        Build an index for the given list of documents.

        Parameters:
            documents (List[Document]): The list of documents to be indexed.

        Returns:
            dict: The built index.
        """
        index = {}
        for doc in documents:
            for key, value in doc.metadata.items():
                if key not in index:
                    index[key] = {}
                if value not in index[key]:
                    index[key][value] = []
                index[key][value].append(doc)
        return index

    def retrieve_metadata(self, search_dict):
        """
        Retrieve documents based on the search criteria provided in search_dict.

        Parameters:
            search_dict (dict): Dictionary specifying the search criteria.
                                It can contain "AND" or "OR" operators for
                                complex queries.

        Returns:
            List[Document]: List of documents that match the search criteria.
        """
        if "AND" in search_dict:
            return self._handle_and(search_dict["AND"])
        elif "OR" in search_dict:
            return self._handle_or(search_dict["OR"])
        else:
            return self._handle_single(search_dict)

    def _handle_and(self, search_dicts):
        results = [self.retrieve_metadata(sd) for sd in search_dicts]
        if results:
            intersection = set.intersection(
                *[set(map(self._hash_doc, r)) for r in results]
            )
            return [self._unhash_doc(h) for h in intersection]
        else:
            return []

    def _handle_or(self, search_dicts):
        results = [self.retrieve_metadata(sd) for sd in search_dicts]
        union = set.union(*[set(map(self._hash_doc, r)) for r in results])
        return [self._unhash_doc(h) for h in union]

    def _handle_single(self, search_dict):
        unions = []
        for key, query in search_dict.items():
            operator, value = query
            union = set()
            if operator == IndexerOperator.EQ:
                if key in self.index and value in self.index[key]:
                    union.update(map(self._hash_doc, self.index[key][value]))
            else:
                if key in self.index:
                    for k, v in self.index[key].items():
                        if (
                            (operator == IndexerOperator.GT and k > value)
                            or (operator == IndexerOperator.GTE and k >= value)
                            or (operator == IndexerOperator.LT and k < value)
                            or (operator == IndexerOperator.LTE and k <= value)
                        ):
                            union.update(map(self._hash_doc, v))
            if union:
                unions.append(union)

        if unions:
            intersection = set.intersection(*unions)
            return [self._unhash_doc(h) for h in intersection]
        else:
            return []

    def _hash_doc(self, doc):
        return (doc.page_content, frozenset(doc.metadata.items()))

    def _unhash_doc(self, hashed_doc):
        page_content, metadata = hashed_doc
        return Document(page_content=page_content, metadata=dict(metadata))


def _get_chat_history(chat_history: List[ChatTurnType]) -> str:
    buffer = ""
    for dialogue_turn in chat_history:
        if isinstance(dialogue_turn, BaseMessage):
            role_prefix = _ROLE_MAP.get(dialogue_turn.type, f"{dialogue_turn.type}: ")
            buffer += f"\n{role_prefix}{dialogue_turn.content}"
        elif isinstance(dialogue_turn, tuple):
            human = "Human: " + dialogue_turn[0]
            ai = "Assistant: " + dialogue_turn[1]
            buffer += "\n" + "\n".join([human, ai])
        else:
            raise ValueError(
                f"Unsupported chat history format: {type(dialogue_turn)}."
                f" Full chat history: {chat_history} "
            )
    return buffer


def _get_standalone_questions_list(
    standalone_questions_str: str, original_question: str
) -> List[str]:
    pattern = r"\d+\.\s(.*?)(?=\n\d+\.|\n|$)"

    matches = [
        match.group(1) for match in re.finditer(pattern, standalone_questions_str)
    ]
    if matches:
        return matches

    match = re.search(
        r"(?i)standalone[^\n]*:[^\n](.*)", standalone_questions_str, re.DOTALL
    )
    sentence_source = match.group(1).strip() if match else standalone_questions_str
    sentences = sentence_source.split("\n")

    return [
        re.sub(
            r"^\((\d+)\)\.? ?|^\d+\.? ?\)?|^(\d+)\) ?|^(\d+)\) ?|^[Qq]uery \d+: ?|^[Qq]uery: ?",
            "",
            sentence.strip(),
        )
        for sentence in sentences
        if sentence.strip()
    ]