Create custom.py

custom.py

@@ -0,0 +1,148 @@
+import os
+from langdetect import detect
+import torch.multiprocessing as mp
+
+from colbert import Indexer, Searcher
+from colbert.infra import ColBERTConfig, Run
+from colbert.utils.utils import print_message
+from colbert.data.collection import Collection
+from colbert.modeling.checkpoint import Checkpoint
+from colbert.indexing.index_saver import IndexSaver
+from colbert.search.index_storage import IndexScorer
+from colbert.infra.launcher import Launcher, print_memory_stats
+from colbert.indexing.collection_encoder import CollectionEncoder
+from colbert.indexing.collection_indexer import CollectionIndexer
+
+
+MMARCO_LANGUAGES = {
+    'ar': ('arabic', 'ar_AR'),
+    'de': ('german', 'de_DE'),
+    'en': ('english', 'en_XX'),
+    'es': ('spanish', 'es_XX'),
+    'fr': ('french', 'fr_XX'),
+    'hi': ('hindi', 'hi_IN'),
+    'id': ('indonesian', 'id_ID'),
+    'it': ('italian', 'it_IT'),
+    'ja': ('japanese', 'ja_XX'),
+    'nl': ('dutch', 'nl_XX'),
+    'pt': ('portuguese', 'pt_XX'),
+    'ru': ('russian', 'ru_RU'),
+    'vi': ('vietnamese', 'vi_VN'),
+    'zh': ('chinese', 'zh_CN'),
+}
+MRTYDI_LANGUAGES = {
+    'ar': ('arabic', 'ar_AR'),
+    'bn': ('bengali', 'bn_IN'),
+    'en': ('english', 'en_XX'),
+    'fi': ('finnish', 'fi_FI'),
+    'id': ('indonesian', 'id_ID'),
+    'ja': ('japanese', 'ja_XX'),
+    'ko': ('korean', 'ko_KR'),
+    'ru': ('russian', 'ru_RU'),
+    'sw': ('swahili', 'sw_KE'),
+    'te': ('telugu', 'te_IN'),
+    'th': ('thai', 'th_TH'),
+}
+MIRACL_LANGUAGES = {
+    'ar': ('arabic', 'ar_AR'),
+    'bn': ('bengali', 'bn_IN'),
+    'en': ('english', 'en_XX'),
+    'es': ('spanish', 'es_XX'),
+    'fa': ('persian', 'fa_IR'),
+    'fi': ('finnish', 'fi_FI'),
+    'fr': ('french', 'fr_XX'),
+    'hi': ('hindi', 'hi_IN'),
+    'id': ('indonesian', 'id_ID'),
+    'ja': ('japanese', 'ja_XX'),
+    'ko': ('korean', 'ko_KR'),
+    'ru': ('russian', 'ru_RU'),
+    'sw': ('swahili', 'sw_KE'),
+    'te': ('telugu', 'te_IN'),
+    'th': ('thai', 'th_TH'),
+    'zh': ('chinese', 'zh_CN'),
+}
+ALL_LANGUAGES = {**MMARCO_LANGUAGES, **MRTYDI_LANGUAGES, **MIRACL_LANGUAGES}
+
+
+def set_xmod_language(model, lang:str):
+    """
+    Set the default language code for the model. This is used when the language is not specified in the input.
+    Source: https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/xmod/modeling_xmod.py#L687
+    """
+    lang = lang.split('-')[0]
+    if (value := ALL_LANGUAGES.get(lang)) is not None:
+        model.set_default_language(value[1])
+    else:
+        raise KeyError(f"Language {lang} not supported.")
+
+#-----------------------------------------------------------------------------------------------------------------#
+#                                               INDEXER
+#-----------------------------------------------------------------------------------------------------------------#
+class CustomIndexer(Indexer):
+    def __launch(self, collection):
+        manager = mp.Manager()
+        shared_lists = [manager.list() for _ in range(self.config.nranks)]
+        shared_queues = [manager.Queue(maxsize=1) for _ in range(self.config.nranks)]
+        launcher = Launcher(custom_encode)
+        launcher.launch(self.config, collection, shared_lists, shared_queues, self.verbose)
+
+def custom_encode(config, collection, shared_lists, shared_queues, verbose: int = 3):
+    encoder = CustomCollectionIndexer(config=config, collection=collection, verbose=verbose)
+    encoder.run(shared_lists)
+
+class CustomCollectionIndexer(CollectionIndexer):
+    def __init__(self, config: ColBERTConfig, collection, verbose=2):
+        self.verbose = verbose
+        self.config = config
+        self.rank, self.nranks = self.config.rank, self.config.nranks
+        self.use_gpu = self.config.total_visible_gpus > 0
+        if self.config.rank == 0 and self.verbose > 1:
+            self.config.help()
+        self.collection = Collection.cast(collection)
+        self.checkpoint = Checkpoint(self.config.checkpoint, colbert_config=self.config)
+        if self.checkpoint.bert.__class__.__name__.lower().startswith("xmod"):
+            language = detect(self.collection.__getitem__(0))
+            Run().print_main(f"#> Setting X-MOD language adapters to {language}.")
+            set_xmod_language(self.checkpoint.bert, lang=language)
+        if self.use_gpu:
+            self.checkpoint = self.checkpoint.cuda()
+        self.encoder = CollectionEncoder(config, self.checkpoint)
+        self.saver = IndexSaver(config)
+        print_memory_stats(f'RANK:{self.rank}')
+
+#-----------------------------------------------------------------------------------------------------------------#
+#                                               SEARCHER
+#-----------------------------------------------------------------------------------------------------------------#
+class CustomSearcher(Searcher):
+    def __init__(self, index, checkpoint=None, collection=None, config=None, index_root=None, verbose:int = 3):
+        self.verbose = verbose
+        if self.verbose > 1:
+            print_memory_stats()
+
+        initial_config = ColBERTConfig.from_existing(config, Run().config)
+
+        default_index_root = initial_config.index_root_
+        index_root = index_root if index_root else default_index_root
+        self.index = os.path.join(index_root, index)
+        self.index_config = ColBERTConfig.load_from_index(self.index)
+
+        self.checkpoint = checkpoint or self.index_config.checkpoint
+        self.checkpoint_config = ColBERTConfig.load_from_checkpoint(self.checkpoint)
+        self.config = ColBERTConfig.from_existing(self.checkpoint_config, self.index_config, initial_config)
+
+        self.collection = Collection.cast(collection or self.config.collection)
+        self.configure(checkpoint=self.checkpoint, collection=self.collection)
+
+        self.checkpoint = Checkpoint(self.checkpoint, colbert_config=self.config, verbose=self.verbose)
+        if self.checkpoint.bert.__class__.__name__.lower().startswith("xmod"):
+            language = detect(self.collection.__getitem__(0))
+            print_message(f"#> Setting X-MOD language adapters to {language}.")
+            set_xmod_language(self.checkpoint.bert, lang=language)
+        use_gpu = self.config.total_visible_gpus > 0
+        if use_gpu:
+            self.checkpoint = self.checkpoint.cuda()
+        load_index_with_mmap = self.config.load_index_with_mmap
+        if load_index_with_mmap and use_gpu:
+            raise ValueError(f"Memory-mapped index can only be used with CPU!")
+        self.ranker = IndexScorer(self.index, use_gpu, load_index_with_mmap)
+        print_memory_stats()