llm-ie 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

llm_ie/extractors.py

@@ -1,11 +1,14 @@
 import abc
 import re
+import copy
 import json
 import json_repair
 import inspect
 import importlib.resources
 import warnings
 import itertools
+import asyncio
+import nest_asyncio
 from typing import Set, List, Dict, Tuple, Union, Callable
 from llm_ie.data_types import LLMInformationExtractionFrame, LLMInformationExtractionDocument
 from llm_ie.engines import InferenceEngine
@@ -19,7 +22,7 @@ class Extractor:
         This is the abstract class for (frame and relation) extractors.
         Input LLM inference engine, system prompt (optional), prompt template (with instruction, few-shot examples).
 
-        Parameters
+        Parameters:
         ----------
         inference_engine : InferenceEngine
             the LLM inferencing engine object. Must implements the chat() method.
@@ -38,16 +41,20 @@ class Extractor:
         """
         This method returns the pre-defined prompt guideline for the extractor from the package asset.
         """
+        # Check if the prompt guide is available
         file_path = importlib.resources.files('llm_ie.asset.prompt_guide').joinpath(f"{cls.__name__}_prompt_guide.txt")
-        with open(file_path, 'r', encoding="utf-8") as f:
-            return f.read()
-
+        try:     
+            with open(file_path, 'r', encoding="utf-8") as f:
+                return f.read()
+        except FileNotFoundError:
+            warnings.warn(f"Prompt guide for {cls.__name__} is not available. Is it a customed extractor?", UserWarning)
+            return None
 
     def _get_user_prompt(self, text_content:Union[str, Dict[str,str]]) -> str:
         """
         This method applies text_content to prompt_template and returns a prompt.
 
-        Parameters
+        Parameters:
         ----------
         text_content : Union[str, Dict[str,str]]
             the input text content to put in prompt template. 
@@ -133,7 +140,7 @@ class FrameExtractor(Extractor):
         This is the abstract class for frame extraction.
         Input LLM inference engine, system prompt (optional), prompt template (with instruction, few-shot examples).
 
-        Parameters
+        Parameters:
         ----------
         inference_engine : InferenceEngine
             the LLM inferencing engine object. Must implements the chat() method.
@@ -175,7 +182,7 @@ class FrameExtractor(Extractor):
         the substring must start with the same word token as the pattern. This is due to the observation that 
         LLM often generate the first few words consistently. 
 
-        Parameters
+        Parameters:
         ----------
         text : str
             the input text.
@@ -219,7 +226,7 @@ class FrameExtractor(Extractor):
         outputs a list of spans (2-tuple) for each entity.
         Entities that are not found in the text will be None from output.
 
-        Parameters
+        Parameters:
         ----------
         text : str
             text that contains entities
@@ -241,7 +248,10 @@ class FrameExtractor(Extractor):
 
         # Match entities
         entity_spans = []
-        for entity in entities:          
+        for entity in entities:   
+            if not isinstance(entity, str):  
+                entity_spans.append(None) 
+                continue    
             if not case_sensitive:
                 entity = entity.lower()
 
@@ -322,7 +332,7 @@ class BasicFrameExtractor(FrameExtractor):
         Input system prompt (optional), prompt template (with instruction, few-shot examples), 
         and specify a LLM.
 
-        Parameters
+        Parameters:
         ----------
         inference_engine : InferenceEngine
             the LLM inferencing engine object. Must implements the chat() method.
@@ -555,18 +565,18 @@ class SentenceFrameExtractor(FrameExtractor):
     from nltk.tokenize.punkt import PunktSentenceTokenizer
     def __init__(self, inference_engine:InferenceEngine, prompt_template:str, system_prompt:str=None, **kwrs):
         """
-        This class performs sentence-based information extraction.
-        A simulated chat follows this process:
+        This class performs sentence-by-sentence information extraction.
+        The process is as follows:
             1. system prompt (optional)
-            2. user instructions (schema, background, full text, few-shot example...)
-            3. user input first sentence
-            4. assistant extract outputs
+            2. user prompt with instructions (schema, background, full text, few-shot example...)
+            3. feed a sentence (start with first sentence)
+            4. LLM extract entities and attributes from the sentence
             5. repeat #3 and #4
 
         Input system prompt (optional), prompt template (with user instructions), 
         and specify a LLM.
 
-        Parameters
+        Parameters:
         ----------
         inference_engine : InferenceEngine
             the LLM inferencing engine object. Must implements the chat() method.
@@ -583,7 +593,7 @@ class SentenceFrameExtractor(FrameExtractor):
         This method sentence tokenize the input text into a list of sentences 
         as dict of {start, end, sentence_text}
 
-        Parameters
+        Parameters:
         ----------
         text : str
             text to sentence tokenize.
@@ -674,10 +684,80 @@ class SentenceFrameExtractor(FrameExtractor):
         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]]:
+        """
+        The asynchronous version of the extract() method.
+
+        Parameters:
+        ----------
+        text_content : Union[str, Dict[str,str]]
+            the input text content to put in prompt template. 
+            If str, the prompt template must has only 1 placeholder {{<placeholder name>}}, regardless of placeholder name.
+            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 should generate. 
+        document_key : str, Optional
+            specify the key in text_content where document text is. 
+            If text_content is str, this parameter will be ignored.
+        temperature : float, Optional
+            the temperature for token sampling.
+        concurrent_batch_size : int, Optional
+            the number of sentences to process in concurrent.
+        """
+        # 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.")
+
+        # define output
+        output = []
+        # sentence tokenization
+        if isinstance(text_content, str):
+            sentences = self._get_sentences(text_content)
+        elif isinstance(text_content, dict):
+            sentences = self._get_sentences(text_content[document_key])
+        # construct chat messages
+        base_messages = []
+        if self.system_prompt:
+            base_messages.append({'role': 'system', 'content': self.system_prompt})
+
+        base_messages.append({'role': 'user', 'content': self._get_user_prompt(text_content)})
+        base_messages.append({'role': 'assistant', 'content': 'Sure, please start with the first sentence.'})
+
+        # generate sentence by sentence
+        tasks = []
+        for i in range(0, len(sentences), concurrent_batch_size):
+            batch = sentences[i:i + concurrent_batch_size]
+            for sent in batch:
+                messages = copy.deepcopy(base_messages)
+                messages.append({'role': 'user', 'content': sent['sentence_text']})
+                task = asyncio.create_task(
+                    self.inference_engine.chat_async(
+                                messages=messages, 
+                                max_new_tokens=max_new_tokens, 
+                                temperature=temperature,
+                                **kwrs
+                            )
+                )
+                tasks.append(task)
+
+            # 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})
+        return output
+    
+
     def extract_frames(self, text_content:Union[str, Dict[str,str]], entity_key:str, max_new_tokens:int=512, 
-                       document_key:str=None, multi_turn:bool=False, temperature:float=0.0, 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]:
+                            document_key:str=None, multi_turn:bool=False, 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]:
         """
         This method inputs a text and outputs a list of LLMInformationExtractionFrame
         It use the extract() method and post-process outputs into frames.
@@ -705,6 +785,10 @@ class SentenceFrameExtractor(FrameExtractor):
             the temperature for token sampling.
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
+        concurrent : bool, Optional
+            if True, the sentences will be extracted in concurrent.
+        concurrent_batch_size : int, Optional
+            the number of sentences to process in concurrent. Only used when `concurrent` is True.
         case_sensitive : bool, Optional
             if True, entity text matching will be case-sensitive.
         fuzzy_match : bool, Optional
@@ -718,15 +802,30 @@ class SentenceFrameExtractor(FrameExtractor):
         Return : str
             a list of frames.
         """
-        llm_output_sentence = self.extract(text_content=text_content, 
-                                           max_new_tokens=max_new_tokens, 
-                                           document_key=document_key,
-                                           multi_turn=multi_turn, 
-                                           temperature=temperature, 
-                                           stream=stream,
-                                           **kwrs)
+        if concurrent:
+            if stream:
+                warnings.warn("stream=True is not supported in concurrent mode.", RuntimeWarning)
+            if multi_turn:
+                warnings.warn("multi_turn=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)
+                                        )
+        else:
+            llm_output_sentences = self.extract(text_content=text_content, 
+                                            max_new_tokens=max_new_tokens, 
+                                            document_key=document_key,
+                                            multi_turn=multi_turn, 
+                                            temperature=temperature, 
+                                            stream=stream,
+                                            **kwrs)
         frame_list = []
-        for sent in llm_output_sentence:
+        for sent in llm_output_sentences:
             entity_json = []
             for entity in self._extract_json(gen_text=sent['gen_text']):
                 if entity_key in entity:
@@ -891,6 +990,121 @@ class SentenceReviewFrameExtractor(SentenceFrameExtractor):
                             'sentence_text': sent['sentence_text'],
                             'gen_text': gen_text})
         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]]:
+        """
+        The asynchronous version of the extract() method.
+
+        Parameters:
+        ----------
+        text_content : Union[str, Dict[str,str]]
+            the input text content to put in prompt template. 
+            If str, the prompt template must has only 1 placeholder {{<placeholder name>}}, regardless of placeholder name.
+            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 should generate. 
+        document_key : str, Optional
+            specify the key in text_content where document text is. 
+            If text_content is str, this parameter will be ignored.
+        temperature : float, Optional
+            the temperature for token sampling.
+        concurrent_batch_size : int, Optional
+            the number of sentences to process in concurrent.
+
+        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'):
+            raise NotImplementedError(f"{self.inference_engine.__class__.__name__} does not have chat_async() method.")
+        
+        # define output
+        output = []
+        # sentence tokenization
+        if isinstance(text_content, str):
+            sentences = self._get_sentences(text_content)
+        elif isinstance(text_content, dict):
+            sentences = self._get_sentences(text_content[document_key])
+        # construct chat messages
+        base_messages = []
+        if self.system_prompt:
+            base_messages.append({'role': 'system', 'content': self.system_prompt})
+
+        base_messages.append({'role': 'user', 'content': self._get_user_prompt(text_content)})
+        base_messages.append({'role': 'assistant', 'content': 'Sure, please start with the first sentence.'})
+
+        # generate initial outputs sentence by sentence
+        initials = []
+        tasks = []
+        message_list = []
+        for i in range(0, len(sentences), concurrent_batch_size):
+            batch = sentences[i:i + concurrent_batch_size]
+            for sent in batch:
+                messages = copy.deepcopy(base_messages)
+                messages.append({'role': 'user', 'content': sent['sentence_text']})
+                message_list.append(messages)
+                task = asyncio.create_task(
+                    self.inference_engine.chat_async(
+                                messages=messages, 
+                                max_new_tokens=max_new_tokens, 
+                                temperature=temperature,
+                                **kwrs
+                            )
+                )
+                tasks.append(task)
+
+        # Wait until the batch is done, collect results and move on to next batch
+        responses = await asyncio.gather(*tasks)
+        # Collect initials
+        for gen_text, sent, message in zip(responses, sentences, message_list):
+            initials.append({'sentence_start': sent['start'],
+                            'sentence_end': sent['end'],
+                            'sentence_text': sent['sentence_text'],
+                            'gen_text': gen_text,
+                            'messages': message})
+            
+        # Review
+        reviews = []
+        tasks = []
+        for i in range(0, len(initials), concurrent_batch_size):
+            batch = initials[i:i + concurrent_batch_size]
+            for init in batch:
+                messages = init["messages"]
+                initial = init["gen_text"]
+                messages.append({'role': 'assistant', 'content': initial})
+                messages.append({'role': 'user', 'content': self.review_prompt})
+                task = asyncio.create_task(
+                                self.inference_engine.chat_async(
+                                messages=messages, 
+                                max_new_tokens=max_new_tokens, 
+                                temperature=temperature,
+                                **kwrs
+                                )
+                            )
+                tasks.append(task) 
+            
+            responses = await asyncio.gather(*tasks)
+
+        # Collect 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})
+        return output
         
 
 class SentenceCoTFrameExtractor(SentenceFrameExtractor):
@@ -1130,19 +1344,34 @@ class BinaryRelationExtractor(RelationExtractor):
             self.possible_relation_func = possible_relation_func
 
 
-    def _extract_relation(self, frame_1:LLMInformationExtractionFrame, frame_2:LLMInformationExtractionFrame, 
-                      text:str, buffer_size:int=100, max_new_tokens:int=128, temperature:float=0.0, stream:bool=False, **kwrs) -> bool:
+    def _post_process(self, rel_json:str) -> bool:
+        if len(rel_json) > 0:
+            if "Relation" in rel_json[0]:
+                rel = rel_json[0]["Relation"]
+                if isinstance(rel, bool):
+                    return rel
+                elif isinstance(rel, str) and rel in {"True", "False"}:
+                    return eval(rel)
+                else:
+                    warnings.warn('Extractor output JSON "Relation" key does not have bool or {"True", "False"} as value.' + \
+                                'Following default, relation = False.', RuntimeWarning)
+            else:
+                warnings.warn('Extractor output JSON without "Relation" key. Following default, relation = False.', RuntimeWarning)
+        else:
+            warnings.warn('Extractor did not output a JSON list. Following default, relation = False.', RuntimeWarning)
+        return False
+    
+
+    def extract(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
+                temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict]:
         """
-        This method inputs two frames and a ROI text, extracts the binary relation.
+        This method considers all combinations of two frames. Use the possible_relation_func to filter impossible pairs.
+        Outputs pairs that are related.
 
         Parameters:
         -----------
-        frame_1 : LLMInformationExtractionFrame
-            a frame
-        frame_2 : LLMInformationExtractionFrame
-            the other frame
-        text : str
-            the entire document text
+        doc : LLMInformationExtractionDocument
+            a document with frames.
         buffer_size : int, Optional
             the number of characters before and after the two frames in the ROI text.
         max_new_tokens : str, Optional
@@ -1152,51 +1381,111 @@ class BinaryRelationExtractor(RelationExtractor):
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
 
-        Return : bool
-            a relation indicator
+        Return : List[Dict]
+            a list of dict with {"frame_1_id", "frame_2_id"}.
         """
-        roi_text = self._get_ROI(frame_1, frame_2, text, buffer_size=buffer_size)
-        if stream:
-            print(f"\n\n{Fore.GREEN}ROI text:{Style.RESET_ALL} \n{roi_text}\n")
-            print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
-
-        messages = []
-        if self.system_prompt:
-            messages.append({'role': 'system', 'content': self.system_prompt})
+        pairs = itertools.combinations(doc.frames, 2)
+        output = []
+        for frame_1, frame_2 in pairs:
+            pos_rel = self.possible_relation_func(frame_1, frame_2)
 
-        messages.append({'role': 'user', 'content': self._get_user_prompt(text_content={"roi_text":roi_text, 
-                                                                                        "frame_1": str(frame_1.to_dict()),
-                                                                                        "frame_2": str(frame_2.to_dict())}
-                                                                                        )})
-        response = self.inference_engine.chat(
-                    messages=messages,
-                    max_new_tokens=max_new_tokens,
-                    temperature=temperature,
-                    stream=stream,
-                    **kwrs
-                )
-        
-        rel_json = self._extract_json(response)
-        if len(rel_json) > 0:
-            if "Relation" in rel_json[0]:
-                rel = rel_json[0]["Relation"]
-                if isinstance(rel, bool):
-                    return rel
-                elif isinstance(rel, str) and rel in {"True", "False"}:
-                    return eval(rel)
-                else:
-                    warnings.warn('Extractor output JSON "Relation" key does not have bool or {"True", "False"} as value.' + \
-                                  'Following default, relation = False.', RuntimeWarning)
-            else:
-                warnings.warn('Extractor output JSON without "Relation" key. Following default, relation = False.', RuntimeWarning)
-        else:
-            warnings.warn("Extractor did not output a JSON. Following default, relation = False.", RuntimeWarning)
+            if pos_rel:
+                roi_text = self._get_ROI(frame_1, frame_2, doc.text, buffer_size=buffer_size)
+                if stream:
+                    print(f"\n\n{Fore.GREEN}ROI text:{Style.RESET_ALL} \n{roi_text}\n")
+                    print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
+                messages = []
+                if self.system_prompt:
+                    messages.append({'role': 'system', 'content': self.system_prompt})
+
+                messages.append({'role': 'user', 'content': self._get_user_prompt(text_content={"roi_text":roi_text, 
+                                                                                                "frame_1": str(frame_1.to_dict()),
+                                                                                                "frame_2": str(frame_2.to_dict())}
+                                                                                                )})
+
+                gen_text = self.inference_engine.chat(
+                                messages=messages, 
+                                max_new_tokens=max_new_tokens, 
+                                temperature=temperature,
+                                stream=stream,
+                                **kwrs
+                            )
+                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})
 
-        return False
+        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]:
+        """
+        This is the asynchronous version of the extract() method.
+
+        Parameters:
+        -----------
+        doc : LLMInformationExtractionDocument
+            a document with frames.
+        buffer_size : int, Optional
+            the number of characters before and after the two frames in the ROI text.
+        max_new_tokens : str, Optional
+            the max number of new tokens LLM should generate. 
+        temperature : float, Optional
+            the temperature for token sampling.
+        concurrent_batch_size : int, Optional
+            the number of frame pairs to process in concurrent.
+
+        Return : List[Dict]
+            a list of dict with {"frame_1", "frame_2"}. 
+        """
+        # 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)
+        n_frames = len(doc.frames)
+        num_pairs = (n_frames * (n_frames-1)) // 2
+        rel_pair_list = []
+        tasks = []
+        for i in range(0, num_pairs, concurrent_batch_size):
+            batch = list(itertools.islice(pairs, concurrent_batch_size))
+            for frame_1, frame_2 in batch:
+                pos_rel = self.possible_relation_func(frame_1, frame_2)
     
+                if pos_rel:
+                    rel_pair_list.append({'frame_1_id':frame_1.frame_id, 'frame_2_id':frame_2.frame_id})
+                    roi_text = self._get_ROI(frame_1, frame_2, doc.text, buffer_size=buffer_size)
+                    messages = []
+                    if self.system_prompt:
+                        messages.append({'role': 'system', 'content': self.system_prompt})
+
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(text_content={"roi_text":roi_text, 
+                                                                                                    "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, 
+                            max_new_tokens=max_new_tokens, 
+                            temperature=temperature,
+                            **kwrs
+                        )
+                    )
+                    tasks.append(task)
+
+            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)
+
+        return output
+
+
     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, concurrent:bool=False, concurrent_batch_size:int=32, stream:bool=False, **kwrs) -> List[Dict]:
         """
         This method considers all combinations of two frames. Use the possible_relation_func to filter impossible pairs.
 
@@ -1210,6 +1499,10 @@ class BinaryRelationExtractor(RelationExtractor):
             the max number of new tokens LLM should generate. 
         temperature : float, Optional
             the temperature for token sampling.
+        concurrent: bool, Optional
+            if True, the extraction will be done in concurrent.
+        concurrent_batch_size : int, Optional
+            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. 
 
@@ -1222,19 +1515,26 @@ class BinaryRelationExtractor(RelationExtractor):
         if doc.has_duplicate_frame_ids():
             raise ValueError("All frame_ids in the input document must be unique.")
 
-        pairs = itertools.combinations(doc.frames, 2)
-        rel_pair_list = []
-        for frame_1, frame_2 in pairs:
-            pos_rel = self.possible_relation_func(frame_1, frame_2)
-            if pos_rel:
-                rel = self._extract_relation(frame_1=frame_1, frame_2=frame_2, text=doc.text, buffer_size=buffer_size, 
-                                         max_new_tokens=max_new_tokens, temperature=temperature, stream=stream, **kwrs)
-                if rel:
-                    rel_pair_list.append({'frame_1':frame_1.frame_id, 'frame_2':frame_2.frame_id})
-
-        return rel_pair_list
-
-
+        if concurrent:
+            if stream:
+                warnings.warn("stream=True is not supported in concurrent mode.", RuntimeWarning)
+
+            nest_asyncio.apply() # For Jupyter notebook. Terminal does not need this.
+            return asyncio.run(self.extract_async(doc=doc, 
+                                                  buffer_size=buffer_size, 
+                                                  max_new_tokens=max_new_tokens,
+                                                  temperature=temperature,
+                                                  concurrent_batch_size=concurrent_batch_size, 
+                                                  **kwrs)
+                                )
+        else:
+            return self.extract(doc=doc, 
+                                buffer_size=buffer_size, 
+                                max_new_tokens=max_new_tokens,
+                                temperature=temperature,
+                                stream=stream,
+                                **kwrs)
+            
 
 class MultiClassRelationExtractor(RelationExtractor):
     def __init__(self, inference_engine:InferenceEngine, prompt_template:str, possible_relation_types_func: Callable, 
@@ -1279,22 +1579,41 @@ class MultiClassRelationExtractor(RelationExtractor):
 
             self.possible_relation_types_func = possible_relation_types_func
 
-        
-    def _extract_relation(self, frame_1:LLMInformationExtractionFrame, frame_2:LLMInformationExtractionFrame, 
-                      pos_rel_types:List[str], text:str, buffer_size:int=100, max_new_tokens:int=128, temperature:float=0.0, stream:bool=False, **kwrs) -> str:
+
+    def _post_process(self, rel_json:List[Dict], pos_rel_types:List[str]) -> Union[str, None]:
         """
-        This method inputs two frames and a ROI text, extracts the relation.
+        This method post-processes the extracted relation JSON.
 
         Parameters:
         -----------
-        frame_1 : LLMInformationExtractionFrame
-            a frame
-        frame_2 : LLMInformationExtractionFrame
-            the other frame
+        rel_json : List[Dict]
+            the extracted relation JSON.
         pos_rel_types : List[str]
-            possible relation types.
-        text : str
-            the entire document text
+            possible relation types by the possible_relation_types_func.
+
+        Return : Union[str, None]
+            the relation type (str) or None for no relation.
+        """
+        if len(rel_json) > 0:
+            if "RelationType" in rel_json[0]:
+                if rel_json[0]["RelationType"] in pos_rel_types:
+                    return rel_json[0]["RelationType"]
+            else:
+                warnings.warn('Extractor output JSON without "RelationType" key. Following default, relation = "No Relation".', RuntimeWarning)
+        else:
+            warnings.warn('Extractor did not output a JSON. Following default, relation = "No Relation".', RuntimeWarning)
+        return None
+    
+
+    def extract(self, doc:LLMInformationExtractionDocument, buffer_size:int=100, max_new_tokens:int=128, 
+                temperature:float=0.0, stream:bool=False, **kwrs) -> List[Dict]:
+        """
+        This method considers all combinations of two frames. Use the possible_relation_types_func to filter impossible pairs.
+
+        Parameters:
+        -----------
+        doc : LLMInformationExtractionDocument
+            a document with frames.
         buffer_size : int, Optional
             the number of characters before and after the two frames in the ROI text.
         max_new_tokens : str, Optional
@@ -1304,54 +1623,117 @@ class MultiClassRelationExtractor(RelationExtractor):
         stream : bool, Optional
             if True, LLM generated text will be printed in terminal in real-time. 
 
-        Return : str
-            a relation type 
+        Return : List[Dict]
+            a list of dict with {"frame_1", "frame_2", "relation"} for all frame pairs.
         """
-        roi_text = self._get_ROI(frame_1, frame_2, text, buffer_size=buffer_size)
-        if stream:
-            print(f"\n\n{Fore.GREEN}ROI text:{Style.RESET_ALL} \n{roi_text}\n")
-            print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
+        pairs = itertools.combinations(doc.frames, 2)
+        output = []
+        for frame_1, frame_2 in pairs:
+            pos_rel_types = self.possible_relation_types_func(frame_1, frame_2)
 
-        messages = []
-        if self.system_prompt:
-            messages.append({'role': 'system', 'content': self.system_prompt})
+            if pos_rel_types:
+                roi_text = self._get_ROI(frame_1, frame_2, doc.text, buffer_size=buffer_size)
+                if stream:
+                    print(f"\n\n{Fore.GREEN}ROI text:{Style.RESET_ALL} \n{roi_text}\n")
+                    print(f"{Fore.BLUE}Extraction:{Style.RESET_ALL}")
+                messages = []
+                if self.system_prompt:
+                    messages.append({'role': 'system', 'content': self.system_prompt})
+
+                messages.append({'role': 'user', 'content': self._get_user_prompt(text_content={"roi_text":roi_text, 
+                                                                                                "frame_1": str(frame_1.to_dict()),
+                                                                                                "frame_2": str(frame_2.to_dict()),
+                                                                                                "pos_rel_types":str(pos_rel_types)}
+                                                                                                )})
+
+                gen_text = self.inference_engine.chat(
+                                messages=messages, 
+                                max_new_tokens=max_new_tokens, 
+                                temperature=temperature,
+                                stream=stream,
+                                **kwrs
+                            )
+                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})
 
-        messages.append({'role': 'user', 'content': self._get_user_prompt(text_content={"roi_text":roi_text, 
-                                                                                        "frame_1": str(frame_1.to_dict()),
-                                                                                        "frame_2": str(frame_2.to_dict()),
-                                                                                        "pos_rel_types":str(pos_rel_types)})})
-        response = self.inference_engine.chat(
-                    messages=messages,
-                    max_new_tokens=max_new_tokens,
-                    temperature=temperature,
-                    stream=stream,
-                    **kwrs
-                )
-        
-        rel_json = self._extract_json(response)
-        if len(rel_json) > 0:
-            if "RelationType" in rel_json[0]:
-                rel = rel_json[0]["RelationType"]
-                if rel in pos_rel_types or rel == "No Relation":
-                    return rel_json[0]["RelationType"]
-                else:
-                    warnings.warn(f'Extracted relation type "{rel}", which is not in the return of possible_relation_types_func: {pos_rel_types}.'+ \
-                                  'Following default, relation = "No Relation".', RuntimeWarning)
-            
-            else:
-                warnings.warn('Extractor output JSON without "RelationType" key. Following default, relation = "No Relation".', RuntimeWarning)
+        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]:
+        """
+        This is the asynchronous version of the extract() method.
+
+        Parameters:
+        -----------
+        doc : LLMInformationExtractionDocument
+            a document with frames.
+        buffer_size : int, Optional
+            the number of characters before and after the two frames in the ROI text.
+        max_new_tokens : str, Optional
+            the max number of new tokens LLM should generate. 
+        temperature : float, Optional
+            the temperature for token sampling.
+        concurrent_batch_size : int, Optional
+            the number of frame pairs to process in concurrent.
+
+        Return : List[Dict]
+            a list of dict with {"frame_1", "frame_2", "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.")
         
-        else:
-            warnings.warn('Extractor did not output a JSON. Following default, relation = "No Relation".', RuntimeWarning)
+        pairs = itertools.combinations(doc.frames, 2)
+        n_frames = len(doc.frames)
+        num_pairs = (n_frames * (n_frames-1)) // 2
+        rel_pair_list = []
+        tasks = []
+        for i in range(0, num_pairs, concurrent_batch_size):
+            batch = list(itertools.islice(pairs, concurrent_batch_size))
+            for frame_1, frame_2 in batch:
+                pos_rel_types = self.possible_relation_types_func(frame_1, frame_2)
+    
+                if pos_rel_types:
+                    rel_pair_list.append({'frame_1':frame_1.frame_id, 'frame_2':frame_2.frame_id, 'pos_rel_types':pos_rel_types})
+                    roi_text = self._get_ROI(frame_1, frame_2, doc.text, buffer_size=buffer_size)
+                    messages = []
+                    if self.system_prompt:
+                        messages.append({'role': 'system', 'content': self.system_prompt})
+
+                    messages.append({'role': 'user', 'content': self._get_user_prompt(text_content={"roi_text":roi_text, 
+                                                                                                    "frame_1": str(frame_1.to_dict()),
+                                                                                                    "frame_2": str(frame_2.to_dict()),
+                                                                                                    "pos_rel_types":str(pos_rel_types)}
+                                                                                                    )})
+                    task = asyncio.create_task(
+                        self.inference_engine.chat_async(
+                            messages=messages, 
+                            max_new_tokens=max_new_tokens, 
+                            temperature=temperature,
+                            **kwrs
+                        )
+                    )
+                    tasks.append(task)
+
+            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})
 
-        return "No Relation"
+        return output
 
 
     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, concurrent:bool=False, concurrent_batch_size:int=32, stream:bool=False, **kwrs) -> List[Dict]:
         """
-        This method considers all combinations of two frames. Use the possible_relation_types_func to filter impossible pairs 
-        and to provide possible relation types between two frames. 
+        This method considers all combinations of two frames. Use the possible_relation_types_func to filter impossible pairs.
 
         Parameters:
         -----------
@@ -1363,6 +1745,10 @@ class MultiClassRelationExtractor(RelationExtractor):
             the max number of new tokens LLM should generate. 
         temperature : float, Optional
             the temperature for token sampling.
+        concurrent: bool, Optional
+            if True, the extraction will be done in concurrent.
+        concurrent_batch_size : int, Optional
+            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. 
 
@@ -1375,15 +1761,23 @@ class MultiClassRelationExtractor(RelationExtractor):
         if doc.has_duplicate_frame_ids():
             raise ValueError("All frame_ids in the input document must be unique.")
 
-        pairs = itertools.combinations(doc.frames, 2)
-        rel_pair_list = []
-        for frame_1, frame_2 in pairs:
-            pos_rel_types = self.possible_relation_types_func(frame_1, frame_2)
-            if pos_rel_types:
-                rel = self._extract_relation(frame_1=frame_1, frame_2=frame_2, pos_rel_types=pos_rel_types, text=doc.text, 
-                                         buffer_size=buffer_size, max_new_tokens=max_new_tokens, temperature=temperature, stream=stream, **kwrs)
-            
-                if rel != "No Relation":
-                    rel_pair_list.append({'frame_1':frame_1.frame_id, 'frame_2':frame_2.frame_id, "relation":rel})
-
-        return rel_pair_list
+        if concurrent:
+            if stream:
+                warnings.warn("stream=True is not supported in concurrent mode.", RuntimeWarning)
+
+            nest_asyncio.apply() # For Jupyter notebook. Terminal does not need this.
+            return asyncio.run(self.extract_async(doc=doc, 
+                                                  buffer_size=buffer_size, 
+                                                  max_new_tokens=max_new_tokens,
+                                                  temperature=temperature,
+                                                  concurrent_batch_size=concurrent_batch_size, 
+                                                  **kwrs)
+                                )
+        else:
+            return self.extract(doc=doc, 
+                                buffer_size=buffer_size, 
+                                max_new_tokens=max_new_tokens,
+                                temperature=temperature,
+                                stream=stream,
+                                **kwrs)
+