customized function for llm invoking

Author: HuXiangkun

README.md

@@ -70,15 +70,15 @@ You can first setup a [demo website](./demo/) and then use the web UI to try Ref
 ## 🚀 Quick Start
 
 ### Setup Environment
-First create a python environment using conda or virtualenv. Clone this repo and change path into the root directory. Then install:
+First create a python environment using conda or virtualenv. Then install:
 ```bash
-pip install -e .
+pip install refchecker
 python -m spacy download en_core_web_sm
 ```
 
 Install optional dependencies to use open source extractors (Mistral, Mixtral) or enable acceleration for RepCChecker.
 ```bash
-pip install -e .[open-extractor,repcex]
+pip install refchecker[open-extractor,repcex]
 ```
 
 ### Code Examples

refchecker/checker/checker_base.py

@@ -49,9 +49,7 @@ def check(
         merge_psg: bool = True,
         is_joint: bool = False,
         joint_check_num: int = 5,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         """
@@ -97,9 +95,7 @@ def check(
                 questions=batch_questions,
                 is_joint=True,
                 joint_check_num=joint_check_num,
-                sagemaker_client=sagemaker_client,
-                sagemaker_params=sagemaker_params,
-                sagemaker_get_response_func=sagemaker_get_response_func,
+                custom_llm_api_func=custom_llm_api_func,
                 **kwargs
             )
             if merge_psg:
@@ -139,9 +135,7 @@ def check(
                     responses=[inp[2] for inp in input_flattened],
                     questions=[inp[3] for inp in input_flattened],
                     is_joint=False,
-                    sagemaker_client=sagemaker_client,
-                    sagemaker_params=sagemaker_params,
-                    sagemaker_get_response_func=sagemaker_get_response_func
+                    custom_llm_api_func=custom_llm_api_func,
                 )
 
             ret = [[x] + y for x, y in zip(ret, input_ids)]

refchecker/checker/llm_checker.py

@@ -46,9 +46,7 @@ def _check(
         questions: List[str] = None,
         is_joint: bool = False,
         joint_check_num: int = 5,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         """
@@ -127,9 +125,7 @@ def _check(
                     model=self.model,
                     max_new_tokens=joint_check_num * 10 + 100,
                     api_base=self.api_base,
-                    sagemaker_client=sagemaker_client,
-                    sagemaker_params=sagemaker_params,
-                    sagemaker_get_response_func=sagemaker_get_response_func,
+                    custom_llm_api_func=custom_llm_api_func,
                     **kwargs
                 )
 
@@ -208,9 +204,7 @@ def _check(
                     model=self.model,
                     max_new_tokens=10,
                     api_base=self.api_base,
-                    sagemaker_client=sagemaker_client,
-                    sagemaker_params=sagemaker_params,
-                    sagemaker_get_response_func=sagemaker_get_response_func,
+                    custom_llm_api_func=custom_llm_api_func,
                     **kwargs
                 )
 

refchecker/extractor/extractor_base.py

@@ -16,29 +16,23 @@ def extract(
         batch_responses, 
         batch_questions=None, 
         max_new_tokens=500,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         if self.claim_format == 'triplet':
             result = self.extract_claim_triplets(
                 batch_responses=batch_responses,
                 batch_questions=batch_questions,
                 max_new_tokens=max_new_tokens,
-                sagemaker_client=sagemaker_client,
-                sagemaker_params=sagemaker_params,
-                sagemaker_get_response_func=sagemaker_get_response_func,
+                custom_llm_api_func=custom_llm_api_func,
                 **kwargs
             )
         elif self.claim_format == 'subsentence':
             result = self.extract_subsentence_claims(
                 batch_responses=batch_responses,
                 batch_questions=batch_questions,
                 max_new_tokens=max_new_tokens,
-                sagemaker_client=sagemaker_client,
-                sagemaker_params=sagemaker_params,
-                sagemaker_get_response_func=sagemaker_get_response_func,
+                custom_llm_api_func=custom_llm_api_func,
                 **kwargs
             )
         return result
@@ -48,9 +42,7 @@ def extract_claim_triplets(
         batch_responses,
         batch_questions=None, 
         max_new_tokens=500,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         raise NotImplementedError
@@ -60,9 +52,7 @@ def extract_subsentence_claims(
         batch_responses,
         batch_questions=None,
         max_new_tokens=500,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         raise NotImplementedError

refchecker/extractor/llm_extractor.py

@@ -30,9 +30,7 @@ def extract_subsentence_claims(
         batch_responses, 
         batch_questions=None, 
         max_new_tokens=500,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         """Extract subsentence claims from the response text.
@@ -75,9 +73,7 @@ def extract_subsentence_claims(
                 n_choices=1,
                 max_new_tokens=max_new_tokens,
                 api_base=self.api_base,
-                sagemaker_client=sagemaker_client,
-                sagemaker_params=sagemaker_params,
-                sagemaker_get_response_func=sagemaker_get_response_func,
+                custom_llm_api_func=custom_llm_api_func,
                 **kwargs
             )
 
@@ -103,9 +99,7 @@ def extract_claim_triplets(
         batch_responses, 
         batch_questions=None, 
         max_new_tokens=500,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
     ):
         """Extract KG triplets from the response text.
@@ -151,9 +145,7 @@ def extract_claim_triplets(
                 n_choices=1,
                 max_new_tokens=max_new_tokens,
                 api_base=self.api_base,
-                sagemaker_client=sagemaker_client,
-                sagemaker_params=sagemaker_params,
-                sagemaker_get_response_func=sagemaker_get_response_func,
+                custom_llm_api_func=custom_llm_api_func,
                 **kwargs
             )
 

refchecker/utils.py

@@ -70,9 +70,7 @@ def get_model_batch_response(
         n_choices=1,
         max_new_tokens=500,
         api_base=None,
-        sagemaker_client=None,
-        sagemaker_params=None,
-        sagemaker_get_response_func=None,
+        custom_llm_api_func=None,
         **kwargs
 ):
     """
@@ -99,35 +97,9 @@ def get_model_batch_response(
     """
     if not prompts or len(prompts) == 0:
         raise ValueError("Invalid input.")
-
-    if sagemaker_client is not None:
-        parameters = {
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature
-        }
-        if sagemaker_params is not None:
-            for k, v in sagemaker_params.items():
-                if k in parameters:
-                    parameters[k] = v
-        response_list = []
-        for prompt in prompts:
-            r = sagemaker_client.invoke_endpoint(
-                EndpointName=model,
-                Body=json.dumps(
-                    {
-                        "inputs": prompt,
-                        "parameters": parameters,
-                    }
-                ),
-                ContentType="application/json",
-            )
-            if sagemaker_get_response_func is not None:
-                response = sagemaker_get_response_func(r)
-            else:
-                r = json.loads(r['Body'].read().decode('utf8'))
-                response = r['outputs'][0]
-            response_list.append(response)
-        return response_list
+    
+    if custom_func is not None:
+        return custom_func(prompts)
     else:
         message_list = []
         for prompt in prompts: