customized function to get response for sagemaker

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "refchecker"
-version = "0.2.8"
+version = "0.2.9"
 description = "RefChecker provides automatic checking pipeline for detecting fine-grained hallucinations generated by Large Language Models."
 authors = [
     "Xiangkun Hu <[email protected]>",

refchecker/checker/checker_base.py

@@ -51,6 +51,7 @@ def check(
         joint_check_num: int = 5,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         """
@@ -98,6 +99,7 @@ def check(
                 joint_check_num=joint_check_num,
                 sagemaker_client=sagemaker_client,
                 sagemaker_params=sagemaker_params,
+                sagemaker_get_response_func=sagemaker_get_response_func,
                 **kwargs
             )
             if merge_psg:
@@ -138,7 +140,8 @@ def check(
                     questions=[inp[3] for inp in input_flattened],
                     is_joint=False,
                     sagemaker_client=sagemaker_client,
-                    sagemaker_params=sagemaker_params
+                    sagemaker_params=sagemaker_params,
+                    sagemaker_get_response_func=sagemaker_get_response_func
                 )
 
             ret = [[x] + y for x, y in zip(ret, input_ids)]

refchecker/checker/llm_checker.py

@@ -48,6 +48,7 @@ def _check(
         joint_check_num: int = 5,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         """
@@ -128,6 +129,7 @@ def _check(
                     api_base=self.api_base,
                     sagemaker_client=sagemaker_client,
                     sagemaker_params=sagemaker_params,
+                    sagemaker_get_response_func=sagemaker_get_response_func,
                     **kwargs
                 )
 
@@ -208,6 +210,7 @@ def _check(
                     api_base=self.api_base,
                     sagemaker_client=sagemaker_client,
                     sagemaker_params=sagemaker_params,
+                    sagemaker_get_response_func=sagemaker_get_response_func,
                     **kwargs
                 )
 

refchecker/extractor/extractor_base.py

@@ -18,6 +18,7 @@ def extract(
         max_new_tokens=500,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         if self.claim_format == 'triplet':
@@ -27,6 +28,7 @@ def extract(
                 max_new_tokens=max_new_tokens,
                 sagemaker_client=sagemaker_client,
                 sagemaker_params=sagemaker_params,
+                sagemaker_get_response_func=sagemaker_get_response_func,
                 **kwargs
             )
         elif self.claim_format == 'subsentence':
@@ -36,6 +38,7 @@ def extract(
                 max_new_tokens=max_new_tokens,
                 sagemaker_client=sagemaker_client,
                 sagemaker_params=sagemaker_params,
+                sagemaker_get_response_func=sagemaker_get_response_func,
                 **kwargs
             )
         return result
@@ -47,6 +50,7 @@ def extract_claim_triplets(
         max_new_tokens=500,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         raise NotImplementedError
@@ -58,6 +62,7 @@ def extract_subsentence_claims(
         max_new_tokens=500,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         raise NotImplementedError

refchecker/extractor/llm_extractor.py

@@ -32,6 +32,7 @@ def extract_subsentence_claims(
         max_new_tokens=500,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         """Extract subsentence claims from the response text.
@@ -76,6 +77,7 @@ def extract_subsentence_claims(
                 api_base=self.api_base,
                 sagemaker_client=sagemaker_client,
                 sagemaker_params=sagemaker_params,
+                sagemaker_get_response_func=sagemaker_get_response_func,
                 **kwargs
             )
 
@@ -103,6 +105,7 @@ def extract_claim_triplets(
         max_new_tokens=500,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
     ):
         """Extract KG triplets from the response text.
@@ -150,6 +153,7 @@ def extract_claim_triplets(
                 api_base=self.api_base,
                 sagemaker_client=sagemaker_client,
                 sagemaker_params=sagemaker_params,
+                sagemaker_get_response_func=sagemaker_get_response_func,
                 **kwargs
             )
 

refchecker/utils.py

@@ -72,6 +72,7 @@ def get_model_batch_response(
         api_base=None,
         sagemaker_client=None,
         sagemaker_params=None,
+        sagemaker_get_response_func=None,
         **kwargs
 ):
     """
@@ -120,9 +121,11 @@ def get_model_batch_response(
                 ),
                 ContentType="application/json",
             )
-
-            r = json.loads(r['Body'].read().decode('utf8'))
-            response = r['outputs'][0]
+            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
     else: