ag-nexla commited on
Commit
5d58ea2
·
1 Parent(s): 2336bf5

updated handler

Browse files
Files changed (2) hide show
  1. colbert_configuration.py +1 -1
  2. handler.py +9 -5
colbert_configuration.py CHANGED
@@ -164,7 +164,7 @@ class ResourceSettings:
164
  @dataclass
165
  class DocSettings:
166
  dim: int = DefaultVal(128)
167
- doc_maxlen: int = DefaultVal(220)
168
  mask_punctuation: bool = DefaultVal(True)
169
 
170
 
 
164
  @dataclass
165
  class DocSettings:
166
  dim: int = DefaultVal(128)
167
+ doc_maxlen: int = DefaultVal(250) # Changed from 220 to 250
168
  mask_punctuation: bool = DefaultVal(True)
169
 
170
 
handler.py CHANGED
@@ -22,11 +22,15 @@ class EndpointHandler:
22
  self.tokenizer = AutoTokenizer.from_pretrained(path)
23
  print(f"Tokenizer loaded from: {path}")
24
 
25
- # Load ColBERTConfig to get doc_maxlen for consistent padding
26
- # IMPORTANT: Use load_from_checkpoint to get the *exact* config used for model export.
27
- self.colbert_config = ColBERTConfig.load_from_checkpoint(path)
28
- self.doc_max_length = self.colbert_config.doc_maxlen
29
- print(f"ColBERTConfig doc_maxlen loaded as: {self.doc_max_length}")
 
 
 
 
30
 
31
  # Load the ONNX model
32
  onnx_model_path = os.path.join(path, "model.onnx")
 
22
  self.tokenizer = AutoTokenizer.from_pretrained(path)
23
  print(f"Tokenizer loaded from: {path}")
24
 
25
+ # Use the doc_maxlen that the ONNX model was *actually exported with* (250).
26
+ # This ensures consistency between the handler's tokenizer and the ONNX model's expectation.
27
+ self.doc_max_length = 250
28
+ print(f"Hardcoded doc_maxlen for tokenizer as: {self.doc_max_length}")
29
+
30
+ # NOTE: If you need other colbert_config parameters, you'd load it here,
31
+ # but for doc_max_length, we are explicitly setting it to avoid mismatches.
32
+ # self.colbert_config = ColBERTConfig.load_from_checkpoint(path)
33
+ # self.doc_max_length = self.colbert_config.doc_maxlen
34
 
35
  # Load the ONNX model
36
  onnx_model_path = os.path.join(path, "model.onnx")