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

llm_ie/extractors.py

@@ -288,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
 
@@ -300,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.
@@ -309,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.
@@ -327,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.
@@ -357,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.
 
@@ -373,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.
@@ -390,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, allow_overlap_entities:bool=False, **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.
@@ -430,6 +442,8 @@ class BasicFrameExtractor(FrameExtractor):
         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.
@@ -442,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):
@@ -472,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
         
 
@@ -518,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.
 
@@ -534,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.
@@ -570,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):
@@ -657,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. 
 
@@ -676,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.
@@ -690,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
@@ -701,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
@@ -724,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'],
@@ -731,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.
 
@@ -754,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'):
@@ -769,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
         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 = []
@@ -783,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
@@ -802,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, batch):
+            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
     
 
@@ -819,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,
-                        allow_overlap_entities:bool=False, **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.
@@ -857,6 +938,8 @@ class SentenceFrameExtractor(FrameExtractor):
         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.
@@ -866,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 = []
@@ -908,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
 
 
@@ -963,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. 
 
@@ -982,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.
@@ -996,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
@@ -1007,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
@@ -1033,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})
 
@@ -1050,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.
 
@@ -1077,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.
@@ -1095,6 +1212,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 initial outputs sentence by sentence
         for i in range(0, len(sentences), concurrent_batch_size):
             messages_list = []
@@ -1111,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
@@ -1175,11 +1305,19 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                 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
         
 
@@ -1221,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. 
 
@@ -1240,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.
@@ -1252,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
@@ -1263,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
@@ -1286,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
     
 
@@ -1361,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. 
 
@@ -1377,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.
@@ -1446,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.
@@ -1463,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)
@@ -1495,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.
 
@@ -1517,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"}. 
@@ -1526,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)
     
@@ -1546,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, 
@@ -1555,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.
 
@@ -1588,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.
@@ -1608,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:
@@ -1616,6 +1808,7 @@ class BinaryRelationExtractor(RelationExtractor):
                                 max_new_tokens=max_new_tokens,
                                 temperature=temperature,
                                 stream=stream,
+                                return_messages_log=return_messages_log,
                                 **kwrs)
             
 
@@ -1689,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.
 
@@ -1705,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)
@@ -1736,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.
 
@@ -1761,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)
     
@@ -1800,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.
 
@@ -1834,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.
@@ -1854,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:
@@ -1862,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.6.dist-info → llm_ie-0.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: llm-ie
-Version: 0.4.6
+Version: 0.4.7
 Summary: An LLM-powered tool that transforms everyday language into robust information extraction pipelines.
 License: MIT
 Author: Enshuo (David) Hsu
@@ -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.6.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=xdg1aYN_My9N9HrzKuMniqFCzNCq7E3DKI7Ru9gawFs,89822
+llm_ie/extractors.py,sha256=O4jt0SxiAk5qKcnIhXqNDGP2Uo16w8UHhsz0Eem_qLQ,100374
 llm_ie/prompt_editor.py,sha256=pHRbg_yFZdoV63r3pvf0TsLfgH2EVJvzUQEVDH1Hj0s,9570
-llm_ie-0.4.6.dist-info/METADATA,sha256=wYKoOc0fvKvEAOvHGnMGOGgaKP49rmDUnHhA8vKBrgg,55745
-llm_ie-0.4.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-llm_ie-0.4.6.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,,