llm-ie 0.1.0__py3-none-any.whl

llm_ie/data_types.py

@@ -0,0 +1,167 @@
+from typing import List, Dict
+import yaml
+
+
+class LLMInformationExtractionFrame:
+    def __init__(self, frame_id:str, start:int, end:int, entity_text:str, attr:Dict[str,str]):
+        """
+        This class holds a frame (entity) extracted by LLM. 
+        A frame contains the span (start and end character positions), a entity text, and 
+        a set of attributes. 
+
+        Parameters
+        ----------
+        frame_id : str
+            unique identiifier for the entity
+        start : int
+            entity start character position
+        end : int
+            entity end character position
+        entity_text : str
+            entity string. Should be the exact string by [start:end]
+        attr : Dict[str,str]
+            dict of attributes
+        """
+        assert isinstance(frame_id, str), "frame_id must be a string."
+        self.frame_id = frame_id
+        self.start = start
+        self.end = end
+        self.entity_text = entity_text
+        self.attr = attr.copy()
+
+    def is_equal(self, frame:"LLMInformationExtractionFrame") -> bool:
+        """ 
+        This method checks if an external frame holds the same information as self.
+        This can be used in evaluation against gold standard. 
+        """
+        return self.start == frame.start and self.end == frame.end
+
+    def is_overlap(self, frame:"LLMInformationExtractionFrame") -> bool:
+        """ 
+        This method checks if an external frame overlaps with self.
+        This can be used in evaluation against gold standard. 
+        """
+        if self.end < frame.start or self.start > frame.end:
+            return False
+        return True
+
+    def to_dict(self) -> Dict[str,str]:
+        """
+        This method outputs the frame contents to a dictionary.
+        """
+        return {"frame_id": self.frame_id,
+                "start": self.start,
+                "end": self.end,
+                "entity_text": self.entity_text,
+                "attr": self.attr}
+    
+    @classmethod
+    def from_dict(cls, d: Dict[str,str]) -> "LLMInformationExtractionFrame":
+        """ 
+        This method defines a LLMInformationExtractionFrame from dictionary.
+        """
+        return cls(frame_id=d['frame_id'],
+                    start=d['start'],
+                    end=d['end'],
+                    entity_text=d['entity_text'],
+                    attr=d['attr'])
+
+    def copy(self) -> "LLMInformationExtractionFrame":
+        return LLMInformationExtractionFrame(frame_id=self.frame_id,
+                                            start=self.start,
+                                            end=self.end,
+                                            entity_text=self.entity_text,
+                                            attr=self.attr)
+
+
+class LLMInformationExtractionDocument:
+    def __init__(self, doc_id:str=None, filename:str=None, text:str=None, frames:List[LLMInformationExtractionFrame]=None):
+        """
+        This class holds LLM-extracted frames, handles save/ load.
+
+        Parameters
+        ----------
+        doc_id : str, Optional
+            document ID. Must be a string
+        filename : str, Optional
+            the directory to a yaml file of a saved LLMInformationExtractionDocument
+        text : str, Optional
+            document text
+        frames : List[LLMInformationExtractionFrame], Optional
+            a list of LLMInformationExtractionFrame
+        """
+        assert doc_id or filename, "Either doc_id (create from raw inputs) or filename (create from file) must be provided."
+        # if create object from file
+        if filename:
+            with open(filename) as yaml_file:
+                llm_ie = yaml.safe_load(yaml_file)
+            if 'doc_id' in llm_ie.keys():
+                self.doc_id = llm_ie['doc_id']
+            if 'text' in llm_ie.keys():
+                self.text = llm_ie['text']
+            if 'frames' in llm_ie.keys():
+                self.frames = [LLMInformationExtractionFrame.from_dict(d) for d in llm_ie['frames']]
+
+        # create object from raw inputs
+        else:
+            assert isinstance(doc_id, str), "doc_id must be a string."
+            self.doc_id = doc_id
+            self.text = text
+            self.frames = frames.copy() if frames is not None else []
+
+
+    def has_frame(self) -> bool:
+        return bool(self.frames)
+    
+    def add_frame(self, frame:LLMInformationExtractionFrame, valid_mode:str=None, create_id:bool=False) -> bool:
+        """
+        This method add a new frame to the frames (list).
+
+        Parameters
+        ----------
+        frame : LLMInformationExtractionFrame
+            the new frame to add.
+        valid_mode : str, Optional
+            one of {None, "span", "attr"}
+            if None, no validation will be done, add frame.
+            if "span", if the new frame's span is equal to an existing frame, add will fail.
+            if "attr", if the new frame's span and all attributes is equal to an existing frame, add will fail.
+        create_id : bool, Optional
+            Assign a sequential frame ID.
+        """
+        assert valid_mode in {None, "span", "attr"}, 'valid_mode must be one of {None, "span", "attr"}'
+
+        if valid_mode == "span":
+            for exist_frame in self.frames:
+                if exist_frame.is_equal(frame):
+                    return False
+
+        if valid_mode == "attr":
+            for exist_frame in self.frames:
+                if exist_frame.is_equal(frame) and exist_frame.attr == frame.attr:
+                    return False
+        
+        # Add frame
+        frame_clone = frame.copy()
+        if create_id:
+            frame_clone.doc_id = f"{self.doc_id}_{len(self.frames)}"
+
+        self.frames.append(frame_clone)
+        return True
+
+
+    def __repr__(self, N_top_chars:int=100) -> str:
+        text_to_print = self.text[0:N_top_chars]
+        frame_count = len(self.frames)
+        return ''.join((f'LLMInformationExtractionDocument(doc_id="{self.doc_id}...")\n',
+                        f'text="{text_to_print}...",\n',
+                        f'frames={frame_count}'))
+
+    def save(self, filename:str):
+        with open(filename, 'w') as yaml_file:
+            yaml.safe_dump({'doc_id':self.doc_id, 
+                            'text':self.text, 
+                            'frames':[frame.to_dict() for frame in self.frames]}, 
+                            yaml_file, sort_keys=False)
+            yaml_file.flush()
+            

llm_ie/engines.py

@@ -0,0 +1,166 @@
+import abc
+from typing import List, Dict
+
+class InferenceEngine:
+    @abc.abstractmethod
+    def __init__(self):
+        """
+        This is an abstract class to provide interfaces for LLM inference engines. 
+        Children classes that inherts this class can be used in extrators. Must implement chat() method.
+        """
+        return NotImplemented
+
+
+    @abc.abstractmethod
+    def chat(self, messages:List[Dict[str,str]], max_new_tokens:int=2048, temperature:float=0.0, stream:bool=False, **kwrs) -> str:
+        """
+        This method inputs chat messages and outputs LLM generated text.
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        max_new_tokens : str, Optional
+            the max number of new tokens LLM can generate. 
+        temperature : float, Optional
+            the temperature for token sampling. 
+        stream : bool, Optional
+            if True, LLM generated text will be printed in terminal in real-time. 
+        """
+        return NotImplemented
+
+
+class LlamaCppInferenceEngine(InferenceEngine):
+    from llama_cpp import Llama
+    def __init__(self, repo_id:str, gguf_filename:str, n_ctx:int=4096, n_gpu_layers:int=-1, **kwrs):
+        """
+        The Llama.cpp inference engine.
+
+        Parameters:
+        ----------
+        repo_id : str
+            the exact name as shown on Huggingface repo
+        gguf_filename : str
+            the exact name as shown in Huggingface repo -> Files and versions. 
+            If multiple gguf files are needed, use the first.
+        n_ctx : int, Optional
+            context length that LLM will evaluate. 
+        n_gpu_layers : int, Optional
+            number of layers to offload to GPU. Default is all layers (-1).
+        """
+        super().__init__()
+        self.repo_id = repo_id
+        self.gguf_filename = gguf_filename
+        self.n_ctx = n_ctx
+        self.n_gpu_layers = n_gpu_layers
+
+        self.model = self.Llama.from_pretrained(
+            repo_id=self.repo_id,
+            filename=self.gguf_filename,
+            n_gpu_layers=n_gpu_layers,
+            n_ctx=n_ctx,
+            **kwrs
+        )
+
+    def __del__(self):
+        """
+        When the inference engine is deleted, release memory for model.
+        """
+        del self.model
+
+
+    def chat(self, messages:List[Dict[str,str]], max_new_tokens:int=2048, temperature:float=0.0, stream:bool=False, **kwrs) -> str:
+        """
+        This method inputs chat messages and outputs LLM generated text.
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        max_new_tokens : str, Optional
+            the max number of new tokens LLM can generate. 
+        temperature : float, Optional
+            the temperature for token sampling. 
+        stream : bool, Optional
+            if True, LLM generated text will be printed in terminal in real-time. 
+        """
+        response = self.model.create_chat_completion(
+                    messages=messages,
+                    max_tokens=max_new_tokens, 
+                    temperature=temperature,
+                    stream=stream,
+                    **kwrs
+                )
+
+        if stream:
+            res = ''
+            for chunk in response:
+                out_dict = chunk['choices'][0]['delta']
+                if 'content' in out_dict:
+                    res += out_dict['content']
+                    print(out_dict['content'], end='', flush=True)
+            print('\n')
+            return res
+        
+        return response['choices'][0]['message']['content']
+
+
+class OllamaInferenceEngine(InferenceEngine):
+    import ollama
+    def __init__(self, model_name:str, num_ctx:int=4096, keep_alive:int=300, **kwrs):
+        """
+        The Ollama inference engine.
+
+        Parameters:
+        ----------
+        model_name : str
+            the model name exactly as shown in >> ollama ls
+        num_ctx : int, Optional
+            context length that LLM will evaluate.
+        keep_alive : int, Optional
+            seconds to hold the LLM after the last API call.
+        """
+        self.model_name = model_name
+        self.num_ctx = num_ctx
+        self.keep_alive = keep_alive
+
+    def chat(self, messages:List[Dict[str,str]], max_new_tokens:int=2048, temperature:float=0.0, stream:bool=False, **kwrs) -> str:
+        """
+        This method inputs chat messages and outputs LLM generated text.
+
+        Parameters:
+        ----------
+        messages : List[Dict[str,str]]
+            a list of dict with role and content. role must be one of {"system", "user", "assistant"}
+        max_new_tokens : str, Optional
+            the max number of new tokens LLM can generate. 
+        temperature : float, Optional
+            the temperature for token sampling. 
+        stream : bool, Optional
+            if True, LLM generated text will be printed in terminal in real-time. 
+        """
+        response = self.ollama.chat(
+                            model=self.model_name, 
+                            messages=messages, 
+                            options={'temperature':temperature, 'num_ctx': self.num_ctx, 'num_predict': max_new_tokens, **kwrs},
+                            stream=stream,
+                            keep_alive=self.keep_alive
+                        )
+        if stream:
+            res = ''
+            for chunk in response:
+                res += chunk['message']['content']
+                print(chunk['message']['content'], end='', flush=True)
+            print('\n')
+            return res
+        
+        return response['message']['content']
+    
+    def release_model_memory(self):
+        """
+        Call API again with keep_alive=0 to release memory for the model
+        """
+        self.ollama.chat(model=self.model_name, 
+                        messages=[{'role': 'user', 'content': ''}], 
+                        options={'num_predict': 0},
+                        keep_alive=0)