Method __init__

index/dense/models.py:65–67

(self, model_path)

# TCTColBERT model
class TCTColBERT(nn.Module):
    def __init__(self, model_path) -> None:
        super(TCTColBERT, self).__init__()
        self.model = BertModel.from_pretrained(model_path)

    def forward(self, input_ids, attention_mask, **kwargs):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state