llm-ie 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

llm_ie/extractors.py

@@ -224,7 +224,8 @@ class FrameExtractor(Extractor):
 
 
     def _find_entity_spans(self, text: str, entities: List[str], case_sensitive:bool=False, 
-                           fuzzy_match:bool=True, fuzzy_buffer_size:float=0.2, fuzzy_score_cutoff:float=0.8) -> List[Tuple[int]]:
+                           fuzzy_match:bool=True, fuzzy_buffer_size:float=0.2, fuzzy_score_cutoff:float=0.8,
+                           allow_overlap_entities:bool=False) -> List[Tuple[int]]:
         """
         This function inputs a text and a list of entity text, 
         outputs a list of spans (2-tuple) for each entity.
@@ -245,6 +246,8 @@ class FrameExtractor(Extractor):
         fuzzy_score_cutoff : float, Optional
             the Jaccard score cutoff for fuzzy matching. 
             Matched entity text must have a score higher than this value or a None will be returned.
+        allow_overlap_entities : bool, Optional
+            if True, entities can overlap in the text. 
         """
         # Handle case sensitivity
         if not case_sensitive:
@@ -264,15 +267,17 @@ class FrameExtractor(Extractor):
             if match and entity:
                 start, end = match.span()
                 entity_spans.append((start, end))
-                # Replace the found entity with spaces to avoid finding the same instance again
-                text = text[:start] + ' ' * (end - start) + text[end:]
+                if not allow_overlap_entities:
+                    # Replace the found entity with spaces to avoid finding the same instance again
+                    text = text[:start] + ' ' * (end - start) + text[end:]
             # Fuzzy match
             elif fuzzy_match:
                 closest_substring_span, best_score = self._get_closest_substring(text, entity, buffer_size=fuzzy_buffer_size)
                 if closest_substring_span and best_score >= fuzzy_score_cutoff:
                     entity_spans.append(closest_substring_span)
-                    # Replace the found entity with spaces to avoid finding the same instance again
-                    text = text[:closest_substring_span[0]] + ' ' * (closest_substring_span[1] - closest_substring_span[0]) + text[closest_substring_span[1]:]
+                    if not allow_overlap_entities:
+                        # Replace the found entity with spaces to avoid finding the same instance again
+                        text = text[:closest_substring_span[0]] + ' ' * (closest_substring_span[1] - closest_substring_span[0]) + text[closest_substring_span[1]:]
                 else:
                     entity_spans.append(None)
 
@@ -283,7 +288,7 @@ class FrameExtractor(Extractor):
         return entity_spans
 
     @abc.abstractmethod
-    def extract(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=2048, **kwrs) -> str:
+    def extract(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=2048, return_messages_log:bool=False, **kwrs) -> str:
         """
         This method inputs text content and outputs a string generated by LLM
 
@@ -295,6 +300,8 @@ class FrameExtractor(Extractor):
             If dict, all the keys must be included in the prompt template placeholder {{<placeholder name>}}.
         max_new_tokens : str, Optional
             the max number of new tokens LLM can generate. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -304,7 +311,7 @@ class FrameExtractor(Extractor):
     
     @abc.abstractmethod
     def extract_frames(self, text_content:Union[str, Dict[str,str]], entity_key:str, max_new_tokens:int=2048, 
-                       document_key:str=None, **kwrs) -> List[LLMInformationExtractionFrame]:
+                       document_key:str=None, return_messages_log:bool=False, **kwrs) -> List[LLMInformationExtractionFrame]:
         """
         This method inputs text content and outputs a list of LLMInformationExtractionFrame
         It use the extract() method and post-process outputs into frames.
@@ -322,6 +329,8 @@ class FrameExtractor(Extractor):
         document_key : str, Optional
             specify the key in text_content where document text is. 
             If text_content is str, this parameter will be ignored.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             a list of frames.
@@ -352,7 +361,7 @@ class BasicFrameExtractor(FrameExtractor):
         
 
     def extract(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=2048, 
-                temperature:float=0.0, stream:bool=False, **kwrs) -> str:
+                temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> str:
         """
         This method inputs a text and outputs a string generated by LLM.
 
@@ -368,6 +377,8 @@ class BasicFrameExtractor(FrameExtractor):
             the temperature for token sampling. 
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -385,13 +396,19 @@ class BasicFrameExtractor(FrameExtractor):
                     **kwrs
                 )
         
+        if return_messages_log:
+            messages.append({"role": "assistant", "content": response})
+            messages_log = [messages]
+            return response, messages_log
+        
         return response
     
 
     def extract_frames(self, text_content:Union[str, Dict[str,str]], entity_key:str, max_new_tokens:int=2048, 
                        temperature:float=0.0, document_key:str=None, stream:bool=False,
                        case_sensitive:bool=False, fuzzy_match:bool=True, fuzzy_buffer_size:float=0.2, 
-                       fuzzy_score_cutoff:float=0.8, **kwrs) -> List[LLMInformationExtractionFrame]:
+                       fuzzy_score_cutoff:float=0.8, allow_overlap_entities:bool=False, 
+                       return_messages_log:bool=False, **kwrs) -> List[LLMInformationExtractionFrame]:
         """
         This method inputs a text and outputs a list of LLMInformationExtractionFrame
         It use the extract() method and post-process outputs into frames.
@@ -422,6 +439,11 @@ class BasicFrameExtractor(FrameExtractor):
         fuzzy_score_cutoff : float, Optional
             the Jaccard score cutoff for fuzzy matching. 
             Matched entity text must have a score higher than this value or a None will be returned.
+        allow_overlap_entities : bool, Optional
+            if True, entities can overlap in the text. 
+            Note that this can cause multiple frames to be generated on the same entity span if they have same entity text.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             a list of frames.
@@ -434,11 +456,13 @@ class BasicFrameExtractor(FrameExtractor):
             text = text_content[document_key]
         
         frame_list = []
-        gen_text = self.extract(text_content=text_content, 
-                                max_new_tokens=max_new_tokens, 
-                                temperature=temperature, 
-                                stream=stream,
-                                **kwrs)
+        extraction_results = self.extract(text_content=text_content, 
+                                    max_new_tokens=max_new_tokens, 
+                                    temperature=temperature, 
+                                    stream=stream,
+                                    return_messages_log=return_messages_log,
+                                    **kwrs)
+        gen_text, messages_log = extraction_results if return_messages_log else (extraction_results, None)
 
         entity_json = []
         for entity in self._extract_json(gen_text=gen_text):
@@ -452,7 +476,8 @@ class BasicFrameExtractor(FrameExtractor):
                                         case_sensitive=case_sensitive,
                                         fuzzy_match=fuzzy_match,
                                         fuzzy_buffer_size=fuzzy_buffer_size,
-                                        fuzzy_score_cutoff=fuzzy_score_cutoff)
+                                        fuzzy_score_cutoff=fuzzy_score_cutoff,
+                                        allow_overlap_entities=allow_overlap_entities)
         
         for i, (ent, span) in enumerate(zip(entity_json, spans)):
             if span is not None:
@@ -463,6 +488,10 @@ class BasicFrameExtractor(FrameExtractor):
                             entity_text=text[start:end],
                             attr={k: v for k, v in ent.items() if k != entity_key and v != ""})
                 frame_list.append(frame)
+
+        if return_messages_log:
+            return frame_list, messages_log
+        
         return frame_list
         
 
@@ -509,7 +538,7 @@ class ReviewFrameExtractor(BasicFrameExtractor):
 
 
     def extract(self, text_content:Union[str, Dict[str,str]], 
-                max_new_tokens:int=4096, temperature:float=0.0, stream:bool=False, **kwrs) -> str:
+                max_new_tokens:int=4096, temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> str:
         """
         This method inputs a text and outputs a string generated by LLM.
 
@@ -525,6 +554,8 @@ class ReviewFrameExtractor(BasicFrameExtractor):
             the temperature for token sampling. 
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -561,10 +592,18 @@ class ReviewFrameExtractor(BasicFrameExtractor):
                     )
 
         # Output
+        output_text = ""
         if self.review_mode == "revision":
-            return review
+            output_text = review
         elif self.review_mode == "addition":
-            return initial + '\n' + review
+            output_text = initial + '\n' + review
+
+        if return_messages_log:
+            messages.append({"role": "assistant", "content": review})
+            messages_log = [messages]
+            return output_text, messages_log
+        
+        return output_text
 
 
 class SentenceFrameExtractor(FrameExtractor):
@@ -648,7 +687,7 @@ class SentenceFrameExtractor(FrameExtractor):
     
     
     def extract(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=512, 
-                document_key:str=None, temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict[str,str]]:
+                document_key:str=None, temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict[str,str]]:
         """
         This method inputs a text and outputs a list of outputs per sentence. 
 
@@ -667,6 +706,8 @@ class SentenceFrameExtractor(FrameExtractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -681,6 +722,9 @@ class SentenceFrameExtractor(FrameExtractor):
                 raise ValueError("document_key must be provided when text_content is dict.")
             sentences = self._get_sentences(text_content[document_key])
 
+        if return_messages_log:
+            messages_log = []
+
         # generate sentence by sentence
         for i, sent in enumerate(sentences):
             # construct chat messages
@@ -692,10 +736,20 @@ class SentenceFrameExtractor(FrameExtractor):
             
             if self.context_sentences == 0:
                 # no context, just place sentence of interest
-                messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                if isinstance(text_content, str):
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                else:
+                    sentence_content = text_content.copy()
+                    sentence_content[document_key] = sent['sentence_text']
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(sentence_content)})
             else:
                 # insert context
-                messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                if isinstance(text_content, str):
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                else:
+                    context_content = text_content.copy()
+                    context_content[document_key] = context
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
                 # simulate conversation
                 messages.append({'role': 'assistant', 'content': 'Sure, please provide the sentence of interest.'})
                 # place sentence of interest
@@ -715,6 +769,10 @@ class SentenceFrameExtractor(FrameExtractor):
                             stream=stream,
                             **kwrs
                         )
+            
+            if return_messages_log:
+                messages.append({"role": "assistant", "content": gen_text})
+                messages_log.append(messages)
 
             # add to output
             output.append({'sentence_start': sent['start'],
@@ -722,11 +780,15 @@ class SentenceFrameExtractor(FrameExtractor):
                             'sentence_text': sent['sentence_text'],
                             'gen_text': gen_text})
             
+        if return_messages_log:
+            return output, messages_log
+        
         return output
     
 
     async def extract_async(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=512, 
-                document_key:str=None, temperature:float=0.0, concurrent_batch_size:int=32, **kwrs) -> List[Dict[str,str]]:
+                document_key:str=None, temperature:float=0.0, concurrent_batch_size:int=32, 
+                return_messages_log:bool=False, **kwrs) -> List[Dict[str,str]]:
         """
         The asynchronous version of the extract() method.
 
@@ -745,6 +807,11 @@ class SentenceFrameExtractor(FrameExtractor):
             the temperature for token sampling.
         concurrent_batch_size : int, Optional
             the number of sentences to process in concurrent.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
+
+        Return : str
+            the output from LLM. Need post-processing.
         """
         # Check if self.inference_engine.chat_async() is implemented
         if not hasattr(self.inference_engine, 'chat_async'):
@@ -760,10 +827,14 @@ class SentenceFrameExtractor(FrameExtractor):
                 raise ValueError("document_key must be provided when text_content is dict.")
             sentences = self._get_sentences(text_content[document_key])
 
+        if return_messages_log:
+            messages_log = []
+
         # generate sentence by sentence
-        tasks = []
         for i in range(0, len(sentences), concurrent_batch_size):
+            tasks = []
             batch = sentences[i:i + concurrent_batch_size]
+            batch_messages = []
             for j, sent in enumerate(batch):
                 # construct chat messages
                 messages = []
@@ -774,10 +845,20 @@ class SentenceFrameExtractor(FrameExtractor):
                 
                 if self.context_sentences == 0:
                     # no context, just place sentence of interest
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                    if isinstance(text_content, str):
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                    else:
+                        sentence_content = text_content.copy()
+                        sentence_content[document_key] = sent['sentence_text']
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(sentence_content)})
                 else:
                     # insert context
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                    if isinstance(text_content, str):
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                    else:
+                        context_content = text_content.copy()
+                        context_content[document_key] = context
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
                     # simulate conversation
                     messages.append({'role': 'assistant', 'content': 'Sure, please provide the sentence of interest.'})
                     # place sentence of interest
@@ -793,16 +874,25 @@ class SentenceFrameExtractor(FrameExtractor):
                             )
                 )
                 tasks.append(task)
+                batch_messages.append(messages)
 
             # Wait until the batch is done, collect results and move on to next batch
             responses = await asyncio.gather(*tasks)
 
-        # Collect outputs
-        for gen_text, sent in zip(responses, sentences):
-            output.append({'sentence_start': sent['start'],
-                            'sentence_end': sent['end'],
-                            'sentence_text': sent['sentence_text'],
-                            'gen_text': gen_text})
+            # Collect outputs
+            for gen_text, sent, messages in zip(responses, batch, batch_messages):
+                if return_messages_log:
+                    messages.append({"role": "assistant", "content": gen_text})
+                    messages_log.append(messages)
+
+                output.append({'sentence_start': sent['start'],
+                                'sentence_end': sent['end'],
+                                'sentence_text': sent['sentence_text'],
+                                'gen_text': gen_text})
+        
+        if return_messages_log:
+            return output, messages_log
+        
         return output
     
 
@@ -810,7 +900,7 @@ class SentenceFrameExtractor(FrameExtractor):
                         document_key:str=None, temperature:float=0.0, stream:bool=False, 
                         concurrent:bool=False, concurrent_batch_size:int=32,
                         case_sensitive:bool=False, fuzzy_match:bool=True, fuzzy_buffer_size:float=0.2, fuzzy_score_cutoff:float=0.8,
-                        **kwrs) -> List[LLMInformationExtractionFrame]:
+                        allow_overlap_entities:bool=False, return_messages_log:bool=False, **kwrs) -> List[LLMInformationExtractionFrame]:
         """
         This method inputs a text and outputs a list of LLMInformationExtractionFrame
         It use the extract() method and post-process outputs into frames.
@@ -845,6 +935,11 @@ class SentenceFrameExtractor(FrameExtractor):
         fuzzy_score_cutoff : float, Optional
             the Jaccard score cutoff for fuzzy matching. 
             Matched entity text must have a score higher than this value or a None will be returned.
+        allow_overlap_entities : bool, Optional
+            if True, entities can overlap in the text. 
+            Note that this can cause multiple frames to be generated on the same entity span if they have same entity text.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             a list of frames.
@@ -854,20 +949,25 @@ class SentenceFrameExtractor(FrameExtractor):
                 warnings.warn("stream=True is not supported in concurrent mode.", RuntimeWarning)
 
             nest_asyncio.apply() # For Jupyter notebook. Terminal does not need this.
-            llm_output_sentences = asyncio.run(self.extract_async(text_content=text_content, 
-                                        max_new_tokens=max_new_tokens, 
-                                        document_key=document_key,
-                                        temperature=temperature, 
-                                        concurrent_batch_size=concurrent_batch_size,
-                                        **kwrs)
-                                        )
+            extraction_results = asyncio.run(self.extract_async(text_content=text_content, 
+                                                max_new_tokens=max_new_tokens, 
+                                                document_key=document_key,
+                                                temperature=temperature, 
+                                                concurrent_batch_size=concurrent_batch_size,
+                                                return_messages_log=return_messages_log,
+                                                **kwrs)
+                                            )
         else:
-            llm_output_sentences = self.extract(text_content=text_content, 
+            extraction_results = self.extract(text_content=text_content, 
                                             max_new_tokens=max_new_tokens, 
                                             document_key=document_key,
                                             temperature=temperature, 
                                             stream=stream,
+                                            return_messages_log=return_messages_log,
                                             **kwrs)
+            
+        llm_output_sentences, messages_log = extraction_results if return_messages_log else (extraction_results, None)
+        
         frame_list = []
         for sent in llm_output_sentences:
             entity_json = []
@@ -882,7 +982,8 @@ class SentenceFrameExtractor(FrameExtractor):
                                             case_sensitive=case_sensitive,
                                             fuzzy_match=fuzzy_match,
                                             fuzzy_buffer_size=fuzzy_buffer_size,
-                                            fuzzy_score_cutoff=fuzzy_score_cutoff)
+                                            fuzzy_score_cutoff=fuzzy_score_cutoff,
+                                            allow_overlap_entities=allow_overlap_entities)
             for ent, span in zip(entity_json, spans):
                 if span is not None:
                     start, end = span
@@ -895,6 +996,9 @@ class SentenceFrameExtractor(FrameExtractor):
                                 entity_text=entity_text,
                                 attr={k: v for k, v in ent.items() if k != entity_key and v != ""})
                     frame_list.append(frame)
+
+        if return_messages_log:
+            return frame_list, messages_log
         return frame_list
 
 
@@ -950,7 +1054,7 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
 
 
     def extract(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=512, 
-                document_key:str=None, temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict[str,str]]:
+                document_key:str=None, temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict[str,str]]:
         """
         This method inputs a text and outputs a list of outputs per sentence. 
 
@@ -969,6 +1073,8 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -983,6 +1089,9 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                 raise ValueError("document_key must be provided when text_content is dict.")
             sentences = self._get_sentences(text_content[document_key])
         
+        if return_messages_log:
+            messages_log = []
+
         # generate sentence by sentence
         for i, sent in enumerate(sentences):
             # construct chat messages
@@ -994,10 +1103,20 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
             
             if self.context_sentences == 0:
                 # no context, just place sentence of interest
-                messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                if isinstance(text_content, str):
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                else:
+                    sentence_content = text_content.copy()
+                    sentence_content[document_key] = sent['sentence_text']
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(sentence_content)})
             else:
                 # insert context
-                messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                if isinstance(text_content, str):
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                else:
+                    context_content = text_content.copy()
+                    context_content[document_key] = context
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
                 # simulate conversation
                 messages.append({'role': 'assistant', 'content': 'Sure, please provide the sentence of interest.'})
                 # place sentence of interest
@@ -1020,6 +1139,7 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
             # Review
             if stream:
                 print(f"\n{Fore.YELLOW}Review:{Style.RESET_ALL}")
+
             messages.append({'role': 'assistant', 'content': initial})
             messages.append({'role': 'user', 'content': self.review_prompt})
 
@@ -1037,15 +1157,23 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
             elif self.review_mode == "addition":
                 gen_text = initial + '\n' + review
 
+            if return_messages_log:
+                messages.append({"role": "assistant", "content": review})
+                messages_log.append(messages)
+
             # add to output
             output.append({'sentence_start': sent['start'],
                             'sentence_end': sent['end'],
                             'sentence_text': sent['sentence_text'],
                             'gen_text': gen_text})
+        
+        if return_messages_log:    
+            return output, messages_log
+        
         return output
     
     async def extract_async(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=512, 
-                document_key:str=None, temperature:float=0.0, concurrent_batch_size:int=32, **kwrs) -> List[Dict[str,str]]:
+                document_key:str=None, temperature:float=0.0, concurrent_batch_size:int=32, return_messages_log:bool=False, **kwrs) -> List[Dict[str,str]]:
         """
         The asynchronous version of the extract() method.
 
@@ -1064,6 +1192,8 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
             the temperature for token sampling.
         concurrent_batch_size : int, Optional
             the number of sentences to process in concurrent.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -1082,10 +1212,14 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                 raise ValueError("document_key must be provided when text_content is dict.")
             sentences = self._get_sentences(text_content[document_key])
 
+        if return_messages_log:
+            messages_log = []
+
         # generate initial outputs sentence by sentence
-        tasks = []
-        messages_list = []
         for i in range(0, len(sentences), concurrent_batch_size):
+            messages_list = []
+            init_tasks = []
+            review_tasks = []
             batch = sentences[i:i + concurrent_batch_size]
             for j, sent in enumerate(batch):
                 # construct chat messages
@@ -1097,10 +1231,20 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                 
                 if self.context_sentences == 0:
                     # no context, just place sentence of interest
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                    if isinstance(text_content, str):
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                    else:
+                        sentence_content = text_content.copy()
+                        sentence_content[document_key] = sent['sentence_text']
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(sentence_content)})
                 else:
                     # insert context
-                    messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                    if isinstance(text_content, str):
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                    else:
+                        context_content = text_content.copy()
+                        context_content[document_key] = context
+                        messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
                     # simulate conversation
                     messages.append({'role': 'assistant', 'content': 'Sure, please provide the sentence of interest.'})
                     # place sentence of interest
@@ -1116,24 +1260,21 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                                 **kwrs
                             )
                 )
-                tasks.append(task)
+                init_tasks.append(task)
 
-        # Wait until the batch is done, collect results and move on to next batch
-        responses = await asyncio.gather(*tasks)
-        # Collect initials
-        initials = []
-        for gen_text, sent, messages in zip(responses, sentences, messages_list):
-            initials.append({'sentence_start': sent['start'],
-                            'sentence_end': sent['end'],
-                            'sentence_text': sent['sentence_text'],
-                            'gen_text': gen_text,
-                            'messages': messages})
+            # Wait until the batch is done, collect results and move on to next batch
+            init_responses = await asyncio.gather(*init_tasks)
+            # Collect initials
+            initials = []
+            for gen_text, sent, messages in zip(init_responses, batch, messages_list):
+                initials.append({'sentence_start': sent['start'],
+                                'sentence_end': sent['end'],
+                                'sentence_text': sent['sentence_text'],
+                                'gen_text': gen_text,
+                                'messages': messages})
 
-        # Review
-        tasks = []
-        for i in range(0, len(initials), concurrent_batch_size):
-            batch = initials[i:i + concurrent_batch_size]
-            for init in batch:
+            # Review
+            for init in initials:
                 messages = init["messages"]
                 initial = init["gen_text"]
                 messages.append({'role': 'assistant', 'content': initial})
@@ -1146,29 +1287,37 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                                 **kwrs
                                 )
                             )
-                tasks.append(task) 
+                review_tasks.append(task) 
             
-            responses = await asyncio.gather(*tasks)
-
-        # Collect reviews
-        reviews = []
-        for gen_text, sent in zip(responses, sentences):
-            reviews.append({'sentence_start': sent['start'],
-                            'sentence_end': sent['end'],
-                            'sentence_text': sent['sentence_text'],
-                            'gen_text': gen_text})
-
-        for init, rev in zip(initials, reviews):
-            if self.review_mode == "revision":
-                gen_text = rev['gen_text']
-            elif self.review_mode == "addition":
-                gen_text = init['gen_text'] + '\n' + rev['gen_text']
-
-            # add to output
-            output.append({'sentence_start': init['sentence_start'],
-                            'sentence_end': init['sentence_end'],
-                            'sentence_text': init['sentence_text'],
-                            'gen_text': gen_text})
+            review_responses = await asyncio.gather(*review_tasks)
+
+            # Collect reviews
+            reviews = []
+            for gen_text, sent in zip(review_responses, batch):
+                reviews.append({'sentence_start': sent['start'],
+                                'sentence_end': sent['end'],
+                                'sentence_text': sent['sentence_text'],
+                                'gen_text': gen_text})
+
+            for init, rev in zip(initials, reviews):
+                if self.review_mode == "revision":
+                    gen_text = rev['gen_text']
+                elif self.review_mode == "addition":
+                    gen_text = init['gen_text'] + '\n' + rev['gen_text']
+
+                if return_messages_log:
+                    messages = init["messages"]
+                    messages.append({"role": "assistant", "content": rev['gen_text']})
+                    messages_log.append(messages)
+
+                # add to output
+                output.append({'sentence_start': init['sentence_start'],
+                                'sentence_end': init['sentence_end'],
+                                'sentence_text': init['sentence_text'],
+                                'gen_text': gen_text})
+        
+        if return_messages_log:    
+            return output, messages_log
         return output
         
 
@@ -1210,7 +1359,7 @@ class SentenceCoTFrameExtractor(SentenceFrameExtractor):
         
 
     def extract(self, text_content:Union[str, Dict[str,str]], max_new_tokens:int=512, 
-                document_key:str=None, temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict[str,str]]:
+                document_key:str=None, temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict[str,str]]:
         """
         This method inputs a text and outputs a list of outputs per sentence. 
 
@@ -1229,6 +1378,8 @@ class SentenceCoTFrameExtractor(SentenceFrameExtractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : str
             the output from LLM. Need post-processing.
@@ -1241,6 +1392,9 @@ class SentenceCoTFrameExtractor(SentenceFrameExtractor):
         elif isinstance(text_content, dict):
             sentences = self._get_sentences(text_content[document_key])
 
+        if return_messages_log:
+            messages_log = []
+
         # generate sentence by sentence
         for i, sent in enumerate(sentences):
             # construct chat messages
@@ -1252,10 +1406,20 @@ class SentenceCoTFrameExtractor(SentenceFrameExtractor):
             
             if self.context_sentences == 0:
                 # no context, just place sentence of interest
-                messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                if isinstance(text_content, str):
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(sent['sentence_text'])})
+                else:
+                    sentence_content = text_content.copy()
+                    sentence_content[document_key] = sent['sentence_text']
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(sentence_content)})
             else:
                 # insert context
-                messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                if isinstance(text_content, str):
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(context)})
+                else:
+                    context_content = text_content.copy()
+                    context_content[document_key] = context
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(context_content)})
                 # simulate conversation
                 messages.append({'role': 'assistant', 'content': 'Sure, please provide the sentence of interest.'})
                 # place sentence of interest
@@ -1275,11 +1439,18 @@ class SentenceCoTFrameExtractor(SentenceFrameExtractor):
                             **kwrs
                         )
 
+            if return_messages_log:
+                messages.append({"role": "assistant", "content": gen_text})
+                messages_log.append(messages)
+
             # add to output
             output.append({'sentence_start': sent['start'],
                             'sentence_end': sent['end'],
                             'sentence_text': sent['sentence_text'],
                             'gen_text': gen_text})
+        
+        if return_messages_log:    
+            return output, messages_log
         return output
     
 
@@ -1350,7 +1521,7 @@ class RelationExtractor(Extractor):
 
     @abc.abstractmethod
     def extract_relations(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                         temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict]:
+                         temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This method considers all combinations of two frames. 
 
@@ -1366,6 +1537,8 @@ class RelationExtractor(Extractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
             a list of dict with {"frame_1", "frame_2"} for all relations.
@@ -1435,7 +1608,7 @@ class BinaryRelationExtractor(RelationExtractor):
     
 
     def extract(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict]:
+                temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This method considers all combinations of two frames. Use the possible_relation_func to filter impossible pairs.
         Outputs pairs that are related.
@@ -1452,11 +1625,17 @@ class BinaryRelationExtractor(RelationExtractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
             a list of dict with {"frame_1_id", "frame_2_id"}.
         """
         pairs = itertools.combinations(doc.frames, 2)
+
+        if return_messages_log:
+            messages_log = []
+
         output = []
         for frame_1, frame_2 in pairs:
             pos_rel = self.possible_relation_func(frame_1, frame_2)
@@ -1484,13 +1663,19 @@ class BinaryRelationExtractor(RelationExtractor):
                             )
                 rel_json = self._extract_json(gen_text)
                 if self._post_process(rel_json):
-                    output.append({'frame_1':frame_1.frame_id, 'frame_2':frame_2.frame_id})
+                    output.append({'frame_1_id':frame_1.frame_id, 'frame_2_id':frame_2.frame_id})
 
+                if return_messages_log:
+                    messages.append({"role": "assistant", "content": gen_text})
+                    messages_log.append(messages)
+
+        if return_messages_log:
+            return output, messages_log
         return output
     
     
     async def extract_async(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                            temperature:float=0.0, concurrent_batch_size:int=32, **kwrs) -> List[Dict]:
+                            temperature:float=0.0, concurrent_batch_size:int=32, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This is the asynchronous version of the extract() method.
 
@@ -1506,6 +1691,8 @@ class BinaryRelationExtractor(RelationExtractor):
             the temperature for token sampling.
         concurrent_batch_size : int, Optional
             the number of frame pairs to process in concurrent.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
             a list of dict with {"frame_1", "frame_2"}. 
@@ -1515,12 +1702,17 @@ class BinaryRelationExtractor(RelationExtractor):
             raise NotImplementedError(f"{self.inference_engine.__class__.__name__} does not have chat_async() method.")
         
         pairs = itertools.combinations(doc.frames, 2)
+        if return_messages_log:
+            messages_log = []
+
         n_frames = len(doc.frames)
         num_pairs = (n_frames * (n_frames-1)) // 2
-        rel_pair_list = []
-        tasks = []
+        output = []
         for i in range(0, num_pairs, concurrent_batch_size):
+            rel_pair_list = []
+            tasks = []
             batch = list(itertools.islice(pairs, concurrent_batch_size))
+            batch_messages = []
             for frame_1, frame_2 in batch:
                 pos_rel = self.possible_relation_func(frame_1, frame_2)
     
@@ -1535,6 +1727,7 @@ class BinaryRelationExtractor(RelationExtractor):
                                                                                                     "frame_1": str(frame_1.to_dict()),
                                                                                                     "frame_2": str(frame_2.to_dict())}
                                                                                                     )})
+
                     task = asyncio.create_task(
                         self.inference_engine.chat_async(
                             messages=messages, 
@@ -1544,20 +1737,27 @@ class BinaryRelationExtractor(RelationExtractor):
                         )
                     )
                     tasks.append(task)
+                    batch_messages.append(messages)
 
             responses = await asyncio.gather(*tasks)
 
-        output = []
-        for d, response in zip(rel_pair_list, responses):
-            rel_json = self._extract_json(response)
-            if self._post_process(rel_json):
-                output.append(d)
+            for d, response, messages in zip(rel_pair_list, responses, batch_messages):
+                if return_messages_log:
+                    messages.append({"role": "assistant", "content": response})
+                    messages_log.append(messages)
 
+                rel_json = self._extract_json(response)
+                if self._post_process(rel_json):
+                    output.append(d)
+
+        if return_messages_log:
+            return output, messages_log
         return output
 
 
     def extract_relations(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                         temperature:float=0.0, concurrent:bool=False, concurrent_batch_size:int=32, stream:bool=False, **kwrs) -> List[Dict]:
+                         temperature:float=0.0, concurrent:bool=False, concurrent_batch_size:int=32, 
+                         stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This method considers all combinations of two frames. Use the possible_relation_func to filter impossible pairs.
 
@@ -1577,6 +1777,8 @@ class BinaryRelationExtractor(RelationExtractor):
             the number of frame pairs to process in concurrent.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
             a list of dict with {"frame_1", "frame_2"} for all relations.
@@ -1597,6 +1799,7 @@ class BinaryRelationExtractor(RelationExtractor):
                                                   max_new_tokens=max_new_tokens,
                                                   temperature=temperature,
                                                   concurrent_batch_size=concurrent_batch_size, 
+                                                  return_messages_log=return_messages_log,
                                                   **kwrs)
                                 )
         else:
@@ -1605,6 +1808,7 @@ class BinaryRelationExtractor(RelationExtractor):
                                 max_new_tokens=max_new_tokens,
                                 temperature=temperature,
                                 stream=stream,
+                                return_messages_log=return_messages_log,
                                 **kwrs)
             
 
@@ -1678,7 +1882,7 @@ class MultiClassRelationExtractor(RelationExtractor):
     
 
     def extract(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict]:
+                temperature:float=0.0, stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This method considers all combinations of two frames. Use the possible_relation_types_func to filter impossible pairs.
 
@@ -1694,11 +1898,17 @@ class MultiClassRelationExtractor(RelationExtractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
-            a list of dict with {"frame_1", "frame_2", "relation"} for all frame pairs.
+            a list of dict with {"frame_1_id", "frame_2_id", "relation"} for all frame pairs.
         """
         pairs = itertools.combinations(doc.frames, 2)
+
+        if return_messages_log:
+            messages_log = []
+
         output = []
         for frame_1, frame_2 in pairs:
             pos_rel_types = self.possible_relation_types_func(frame_1, frame_2)
@@ -1725,16 +1935,23 @@ class MultiClassRelationExtractor(RelationExtractor):
                                 stream=stream,
                                 **kwrs
                             )
+                
+                if return_messages_log:
+                    messages.append({"role": "assistant", "content": gen_text})
+                    messages_log.append(messages)
+
                 rel_json = self._extract_json(gen_text)
                 rel = self._post_process(rel_json, pos_rel_types)
                 if rel:
-                    output.append({'frame_1':frame_1.frame_id, 'frame_2':frame_2.frame_id, 'relation':rel})
+                    output.append({'frame_1_id':frame_1.frame_id, 'frame_2_id':frame_2.frame_id, 'relation':rel})
 
+        if return_messages_log:
+            return output, messages_log
         return output   
     
     
     async def extract_async(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                            temperature:float=0.0, concurrent_batch_size:int=32, **kwrs) -> List[Dict]:
+                            temperature:float=0.0, concurrent_batch_size:int=32, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This is the asynchronous version of the extract() method.
 
@@ -1750,21 +1967,28 @@ class MultiClassRelationExtractor(RelationExtractor):
             the temperature for token sampling.
         concurrent_batch_size : int, Optional
             the number of frame pairs to process in concurrent.
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
-            a list of dict with {"frame_1", "frame_2", "relation"} for all frame pairs. 
+            a list of dict with {"frame_1_id", "frame_2_id", "relation"} for all frame pairs. 
         """
         # Check if self.inference_engine.chat_async() is implemented
         if not hasattr(self.inference_engine, 'chat_async'):
             raise NotImplementedError(f"{self.inference_engine.__class__.__name__} does not have chat_async() method.")
         
         pairs = itertools.combinations(doc.frames, 2)
+        if return_messages_log:
+            messages_log = []
+
         n_frames = len(doc.frames)
         num_pairs = (n_frames * (n_frames-1)) // 2
-        rel_pair_list = []
-        tasks = []
+        output = []
         for i in range(0, num_pairs, concurrent_batch_size):
+            rel_pair_list = []
+            tasks = []
             batch = list(itertools.islice(pairs, concurrent_batch_size))
+            batch_messages = []
             for frame_1, frame_2 in batch:
                 pos_rel_types = self.possible_relation_types_func(frame_1, frame_2)
     
@@ -1789,21 +2013,28 @@ class MultiClassRelationExtractor(RelationExtractor):
                         )
                     )
                     tasks.append(task)
+                    batch_messages.append(messages)
 
             responses = await asyncio.gather(*tasks)
 
-        output = []
-        for d, response in zip(rel_pair_list, responses):
-            rel_json = self._extract_json(response)
-            rel = self._post_process(rel_json, d['pos_rel_types'])
-            if rel:
-                output.append({'frame_1':d['frame_1'], 'frame_2':d['frame_2'], 'relation':rel})
+            for d, response, messages in zip(rel_pair_list, responses, batch_messages):
+                if return_messages_log:
+                    messages.append({"role": "assistant", "content": response})
+                    messages_log.append(messages)
+
+                rel_json = self._extract_json(response)
+                rel = self._post_process(rel_json, d['pos_rel_types'])
+                if rel:
+                    output.append({'frame_1_id':d['frame_1'], 'frame_2_id':d['frame_2'], 'relation':rel})
 
+        if return_messages_log:
+            return output, messages_log
         return output
 
 
     def extract_relations(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
-                         temperature:float=0.0, concurrent:bool=False, concurrent_batch_size:int=32, stream:bool=False, **kwrs) -> List[Dict]:
+                         temperature:float=0.0, concurrent:bool=False, concurrent_batch_size:int=32, 
+                         stream:bool=False, return_messages_log:bool=False, **kwrs) -> List[Dict]:
         """
         This method considers all combinations of two frames. Use the possible_relation_types_func to filter impossible pairs.
 
@@ -1823,6 +2054,8 @@ class MultiClassRelationExtractor(RelationExtractor):
             the number of frame pairs to process in concurrent.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        return_messages_log : bool, Optional
+            if True, a list of messages will be returned.
 
         Return : List[Dict]
             a list of dict with {"frame_1", "frame_2", "relation"} for all relations.
@@ -1843,6 +2076,7 @@ class MultiClassRelationExtractor(RelationExtractor):
                                                   max_new_tokens=max_new_tokens,
                                                   temperature=temperature,
                                                   concurrent_batch_size=concurrent_batch_size, 
+                                                  return_messages_log=return_messages_log,
                                                   **kwrs)
                                 )
         else:
@@ -1851,5 +2085,6 @@ class MultiClassRelationExtractor(RelationExtractor):
                                 max_new_tokens=max_new_tokens,
                                 temperature=temperature,
                                 stream=stream,
+                                return_messages_log=return_messages_log,
                                 **kwrs)
             

{llm_ie-0.4.5.dist-info → llm_ie-0.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-ie
-Version: 0.4.5
+Version: 0.4.7
 Summary: An LLM-powered tool that transforms everyday language into robust information extraction pipelines.
 License: MIT
 Author: Enshuo (David) Hsu
@@ -44,7 +44,7 @@ An LLM-powered tool that transforms everyday language into robust information ex
 - [v0.4.5](https://github.com/daviden1013/llm-ie/releases/tag/v0.4.5) (Feb 16, 2025): 
     - Added option to adjust number of context sentences in sentence-based extractors.
     - Added support for OpenAI reasoning models ("o" series).
-
+- [v0.4.6](https://github.com/daviden1013/llm-ie/releases/tag/v0.4.6) (Mar 1, 2025): Allow LLM to output overlapping frames.
 
 ## Table of Contents
 - [Overview](#overview)
@@ -1206,10 +1206,14 @@ We benchmarked the frame and relation extractors on biomedical information extra
 ## Citation
 For more information and benchmarks, please check our paper:
 ```bibtex
-@article{hsu2024llm,
-  title={LLM-IE: A Python Package for Generative Information Extraction with Large Language Models},
+@article{hsu2025llm,
+  title={LLM-IE: a python package for biomedical generative information extraction with large language models},
   author={Hsu, Enshuo and Roberts, Kirk},
-  journal={arXiv preprint arXiv:2411.11779},
-  year={2024}
+  journal={JAMIA open},
+  volume={8},
+  number={2},
+  pages={ooaf012},
+  year={2025},
+  publisher={Oxford University Press}
 }
 ```

{llm_ie-0.4.5.dist-info → llm_ie-0.4.7.dist-info}/RECORD

@@ -16,8 +16,8 @@ llm_ie/asset/prompt_guide/SentenceFrameExtractor_prompt_guide.txt,sha256=YomMhiA
 llm_ie/asset/prompt_guide/SentenceReviewFrameExtractor_prompt_guide.txt,sha256=YomMhiA6BoBIJFOZ1jcoKRqY7_T3PGUsM0t0vd1IuZA,9325
 llm_ie/data_types.py,sha256=_Kt4Er1SMj1jg8U8TCXFJH_64prur-IbFngHKmZgWr8,15717
 llm_ie/engines.py,sha256=nWQzV7mcRCOkJ-U1iP-xrT9dVahVbj-nAjQci4XaRjY,22609
-llm_ie/extractors.py,sha256=JgMtENUXJMXPpPyVBg-QiaSRJdQul6ASnQmAg1nYmiI,88758
+llm_ie/extractors.py,sha256=O4jt0SxiAk5qKcnIhXqNDGP2Uo16w8UHhsz0Eem_qLQ,100374
 llm_ie/prompt_editor.py,sha256=pHRbg_yFZdoV63r3pvf0TsLfgH2EVJvzUQEVDH1Hj0s,9570
-llm_ie-0.4.5.dist-info/METADATA,sha256=Ff_g89MgEFANtaFMilNAS7FXNpdOg7C89XX3uLdR4WE,55621
-llm_ie-0.4.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-llm_ie-0.4.5.dist-info/RECORD,,
+llm_ie-0.4.7.dist-info/METADATA,sha256=B5YmHVvdweau9ZU0ORYPFPeI6s6I7WTcqZ6kuu7Xaeg,55821
+llm_ie-0.4.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+llm_ie-0.4.7.dist-info/RECORD,,