jdwh08s/Autodoc-Lifter

1#####################################################2### DOCUMENT PROCESSOR [Keywords]3#####################################################4### Jonathan Wang5 6# ABOUT: 7# This creates an app to chat with PDFs.8 9# This is the Keywords10# Which creates keywords based on documents.11#####################################################12### TODO Board:13# TODO(Jonathan Wang): Add Maximum marginal relevance to the merger for better keywords.14# TODO(Jonathan Wang): create own version of Rake keywords15 16#####################################################17### PROGRAM SETTINGS18 19 20#####################################################21### PROGRAM IMPORTS22from __future__ import annotations23 24from typing import Any, Callable, Optional25 26# Keywords27# from multi_rake import Rake  # removing because of compile issues and lack of maintainence28import yake29from llama_index.core.bridge.pydantic import Field30from llama_index.core.schema import BaseNode31 32# Own Modules33from metadata_adder import MetadataAdder34 35#####################################################36### SCRIPT37 38def get_keywords(input_text: str) -> str:39    """40    Given a string, get its keywords using RAKE+YAKE w/ Distribution Based Fusion.41 42    Inputs:43        input_text (str): the input text to get keywords from44        # top_k (int): the number of keywords to get45 46    Returns:47        str: A list of the keywords, joined into a string.48    """49    # RAKE50    # kw_extractor = Rake()51    # keywords_rake = kw_extractor.apply(input_text)52    # keywords_rake = dict(keywords_rake)53    # YAKE54    kw_extractor = yake.KeywordExtractor(lan="en", dedupLim=0.9, n=3)55    keywords_yake = kw_extractor.extract_keywords(input_text)56    # reorder scores so that higher is better57    keywords_yake = {keyword[0].lower(): (1 - keyword[1]) for keyword in keywords_yake}58    keywords_yake = dict(59        sorted(keywords_yake.items(), key=lambda x: x[1], reverse=True)  # type hinting YAKE is miserable60        )61 62    # Merge RAKE and YAKE based on scores.63    # keywords_merged = _merge_on_scores(64    #     list(keywords_yake.keys()), 65    #     list(keywords_rake.keys()), 66    #     list(keywords_yake.values()), 67    #     list(keywords_rake.values()), 68    #     a_weight=0.5, 69    #     top_k=top_k70    # )71 72    # return (list(keywords_rake.keys())[:top_k], list(keywords_yake.keys())[:top_k], keywords_merged)73    return ", ".join(keywords_yake)  # kinda regretting forcing this into a string74 75 76class KeywordMetadataAdder(MetadataAdder):77    """Adds keyword metadata to a document.78 79    Args:80        metadata_name: The name of the metadata to add to the document. Defaults to 'keyword_metadata'.81        keywords_function: A function for keywords, given a source string and the number of keywords to get.82    """83 84    keywords_function: Callable[[str, int], str] = Field(85        description="The function to use to extract keywords from the text. Input is string and number of keywords to extract. Ouptut is string of keywords.",86        default=get_keywords,87    )88    num_keywords: int = Field(89        default=5,90        description="The number of keywords to extract from the text. Defaults to 5.",91    )92 93    def __init__(94        self,95        metadata_name: str = "keyword_metadata",96        keywords_function: Callable[[str], str] = get_keywords,97        num_keywords: int = 5,98        **kwargs: Any,99    ) -> None:100        """Init params."""101        super().__init__(metadata_name=metadata_name, keywords_function=keywords_function, num_keywords=num_keywords, **kwargs)  # ah yes i love oop :)102 103    @classmethod104    def class_name(cls) -> str:105        return "KeywordMetadataAdder"106 107    def get_node_metadata(self, node: BaseNode) -> str | None:108        if not hasattr(node, "text") or node.text is None:109            return None110        return self.keywords_function(node.get_content(), self.num_keywords)111