Spaces:
Runtime error
Runtime error
| import os | |
| import torch | |
| from colbert.utils.utils import torch_load_dnn | |
| from transformers import AutoTokenizer | |
| from colbert.modeling.hf_colbert import HF_ColBERT | |
| from colbert.infra.config import ColBERTConfig | |
| class BaseColBERT(torch.nn.Module): | |
| """ | |
| Shallow module that wraps the ColBERT parameters, custom configuration, and underlying tokenizer. | |
| This class provides direct instantiation and saving of the model/colbert_config/tokenizer package. | |
| Like HF, evaluation mode is the default. | |
| """ | |
| def __init__(self, name, colbert_config=None): | |
| super().__init__() | |
| self.name = name | |
| self.colbert_config = ColBERTConfig.from_existing(ColBERTConfig.load_from_checkpoint(name), colbert_config) | |
| self.model = HF_ColBERT.from_pretrained(name, colbert_config=self.colbert_config) | |
| self.raw_tokenizer = AutoTokenizer.from_pretrained(self.model.base) | |
| self.eval() | |
| def device(self): | |
| return self.model.device | |
| def bert(self): | |
| return self.model.bert | |
| def linear(self): | |
| return self.model.linear | |
| def score_scaler(self): | |
| return self.model.score_scaler | |
| def save(self, path): | |
| assert not path.endswith('.dnn'), f"{path}: We reserve *.dnn names for the deprecated checkpoint format." | |
| self.model.save_pretrained(path) | |
| self.raw_tokenizer.save_pretrained(path) | |
| self.colbert_config.save_for_checkpoint(path) | |
| if __name__ == '__main__': | |
| import random | |
| import numpy as np | |
| from colbert.infra.run import Run | |
| from colbert.infra.config import RunConfig | |
| random.seed(12345) | |
| np.random.seed(12345) | |
| torch.manual_seed(12345) | |
| with Run().context(RunConfig(gpus=2)): | |
| m = BaseColBERT('bert-base-uncased', colbert_config=ColBERTConfig(Run().config, doc_maxlen=300, similarity='l2')) | |
| m.colbert_config.help() | |
| print(m.linear.weight) | |
| m.save('/future/u/okhattab/tmp/2021/08/model.deleteme2/') | |
| m2 = BaseColBERT('/future/u/okhattab/tmp/2021/08/model.deleteme2/') | |
| m2.colbert_config.help() | |
| print(m2.linear.weight) | |
| exit() | |
| m = BaseColBERT('/future/u/okhattab/tmp/2021/08/model.deleteme/') | |
| print('BaseColBERT', m.linear.weight) | |
| print('BaseColBERT', m.colbert_config) | |
| exit() | |
| # m = HF_ColBERT.from_pretrained('nreimers/MiniLMv2-L6-H768-distilled-from-BERT-Large') | |
| m = HF_ColBERT.from_pretrained('/future/u/okhattab/tmp/2021/08/model.deleteme/') | |
| print('HF_ColBERT', m.linear.weight) | |
| m.save_pretrained('/future/u/okhattab/tmp/2021/08/model.deleteme/') | |
| # old = OldColBERT.from_pretrained('bert-base-uncased') | |
| # print(old.bert.encoder.layer[10].attention.self.value.weight) | |
| # random.seed(12345) | |
| # np.random.seed(12345) | |
| # torch.manual_seed(12345) | |
| dnn = torch_load_dnn( | |
| "/future/u/okhattab/root/TACL21/experiments/Feb26.NQ/train.py/ColBERT.C3/checkpoints/colbert-60000.dnn") | |
| # base = dnn.get('arguments', {}).get('model', 'bert-base-uncased') | |
| # new = BaseColBERT.from_pretrained('bert-base-uncased', state_dict=dnn['model_state_dict']) | |
| # print(new.bert.encoder.layer[10].attention.self.value.weight) | |
| print(dnn['model_state_dict']['linear.weight']) | |
| # print(dnn['model_state_dict']['bert.encoder.layer.10.attention.self.value.weight']) | |
| # # base_model_prefix | |