Skip to content

Commit

Permalink
Merge pull request #38 from IvVlasov/fix/resource-warning-unclosed-files
Browse files Browse the repository at this point in the history
Fix: ResourceWarning: Enable tracemalloc to get the object allocation…
  • Loading branch information
PrithivirajDamodaran authored Nov 24, 2024
2 parents a7481ad + 96422ab commit 00486e5
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions flashrank/Ranker.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,9 +112,12 @@ def _get_tokenizer(self, max_length: int = 512) -> Tokenizer:
Returns:
Tokenizer: Configured tokenizer for text processing.
"""
config = json.load(open(str(self.model_dir / "config.json")))
tokenizer_config = json.load(open(str(self.model_dir / "tokenizer_config.json")))
tokens_map = json.load(open(str(self.model_dir / "special_tokens_map.json")))
with open(str(self.model_dir / "config.json")) as config_file:
config = json.load(config_file)
with open(str(self.model_dir / "tokenizer_config.json")) as tokenizer_config_file:
tokenizer_config = json.load(tokenizer_config_file)
with open(str(self.model_dir / "special_tokens_map.json")) as tokens_map_file:
tokens_map = json.load(tokens_map_file)
tokenizer = Tokenizer.from_file(str(self.model_dir / "tokenizer.json"))

tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
Expand Down

0 comments on commit 00486e5

Please sign in to comment.