customize chunk_size in score function

handler.py

@@ -55,10 +55,20 @@ class EndpointHandler():
         
         self.tfidf_order = True
         self.num_highlights = 1
+        
+        self.default_chunk_size = 500
+        self.chunk_size = 500
 
 
     def __call__(self, data):
 
+        # this is necessary for setting the chunk size for
+        # retrived docs
+        if 'chunk_size' in data['inputs']:
+            self.chunk_size = int(data['inputs']['chunk_size'])
+        else:
+            self.chunk_size = self.default_chunk_size
+
         claim = data['inputs']['claims'][0]
         ents = extract_entities(claim)
 
@@ -128,9 +138,11 @@ class EndpointHandler():
         retrieved_data = {
             'inputs': {
                 'docs': list(retrieved_docs),
-                'claims': [claim]*len(retrieved_docs)
+                'claims': [claim]*len(retrieved_docs),
+                'chunk_size': self.chunk_size
             }
         }
+
         _, _, used_chunk, support_prob_per_chunk = self.scorer.score(data=retrieved_data)
         end = time()
         num_chunks = len([item for items in used_chunk for item in items])

minicheck_web/inference.py

@@ -28,7 +28,7 @@ def sent_tokenize_with_newlines(text):
 
 
 class Inferencer():
-    def __init__(self, path, chunk_size, max_input_length, batch_size) -> None:
+    def __init__(self, path, max_input_length, batch_size) -> None:
         
         self.path = path
         self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -36,7 +36,9 @@ class Inferencer():
         self.model = AutoModelForSeq2SeqLM.from_pretrained(path).to(self.device)
         self.tokenizer = AutoTokenizer.from_pretrained(path)
 
-        self.chunk_size=500 if chunk_size is None else chunk_size
+        self.default_chunk_size = 500
+        self.chunk_size=500
+
         self.max_input_length=2048 if max_input_length is None else max_input_length
         self.max_output_length = 256
         

minicheck_web/minicheck.py

@@ -9,12 +9,11 @@ import numpy as np
 
 
 class MiniCheck:
-    def __init__(self, path, chunk_size=None, max_input_length=None, batch_size=16) -> None:
+    def __init__(self, path, max_input_length=None, batch_size=16) -> None:
 
         self.model = Inferencer(
             path=path,
             batch_size=batch_size, 
-            chunk_size=chunk_size,
             max_input_length=max_input_length,
         )
 
@@ -30,6 +29,11 @@ class MiniCheck:
         docs = inputs['docs']
         claims = inputs['claims']
 
+        if 'chunk_size' in inputs:
+            self.model.chunk_size = int(inputs['chunk_size'])
+        else:
+            self.model.chunk_size = self.model.default_chunk_size
+
         assert isinstance(docs, list) or isinstance(docs, np.ndarray), f"docs must be a list or np.ndarray"
         assert isinstance(claims, list) or isinstance(claims, np.ndarray), f"claims must be a list or np.ndarray"