import torch from peft import PeftModel from transformers import PreTrainedModel, AutoConfig, T5ForTokenClassification, AutoModel, AutoTokenizer, AutoModelForTokenClassification from .configuration_hhem_v2 import HHEMv2Config class HHEMv2Model(PreTrainedModel): config_class = HHEMv2Config def __init__(self, config): super().__init__(config) # self.t5 = T5ForTokenClassification.from_config( # AutoConfig.from_pretrained(config.foundation) # ) # def populate(self, model): # self.t5 = model # def forward(self, **kwarg): # return self.t5.transformer(**kwarg) class HHEMv2ForSequenceClassification(PreTrainedModel): config_class = HHEMv2Config def __init__(self, config=HHEMv2Config()): super().__init__(config) self.t5 = T5ForTokenClassification( AutoConfig.from_pretrained(config.foundation) ) self.prompt = config.prompt self.tokenzier = AutoTokenizer.from_pretrained(config.foundation) def populate(self, model: AutoModel): """Initiate the model with the pretrained model This method should only be called by Vectara employee who prepares the model for publishing. Users do not need to call this method. """ self.t5 = model # TODO: Figure out how to publish only the adapter yet still able to do end-to-end pulling and inference. # def populate_lora(self, checkpoint: str): # base_model = AutoModelForTokenClassification.from_pretrained(self.config.foundation) # combined_model = PeftModel.from_pretrained(base_model, checkpoint, is_trainable=False) # self.t5 = combined_model def forward(self, **kwargs): return self.t5(**kwargs) def predict(self, text_pairs): tokenizer = self.tokenzier pair_dict = [{'text1': pair[0], 'text2': pair[1]} for pair in text_pairs] inputs = tokenizer( [self.prompt.format(**pair) for pair in pair_dict], return_tensors='pt', padding=True) self.t5.eval() with torch.no_grad(): outputs = self.t5(**inputs) logits = outputs.logits logits = logits[:, 0, :] # tok_cls transformed_probs = torch.softmax(logits, dim=-1) raw_scores = transformed_probs[:, 1] # the probability of class 1 return raw_scores