| from transformers import TokenClassificationPipeline,MistralModel,MistralPreTrainedModel | |
| from transformers.modeling_outputs import TokenClassifierOutput | |
| class BellmanFordTokenClassificationPipeline(TokenClassificationPipeline): | |
| def __init__(self,**kwargs): | |
| import numpy | |
| super().__init__(**kwargs) | |
| x=self.model.config.label2id | |
| y=[k for k in x if not k.startswith("I-")] | |
| self.transition=numpy.full((len(x),len(x)),numpy.nan) | |
| for k,v in x.items(): | |
| for j in ["I-"+k[2:]] if k.startswith("B-") else [k]+y if k.startswith("I-") else y: | |
| self.transition[v,x[j]]=0 | |
| def check_model_type(self,supported_models): | |
| pass | |
| def postprocess(self,model_outputs,**kwargs): | |
| import numpy | |
| if "logits" not in model_outputs: | |
| return self.postprocess(model_outputs[0],**kwargs) | |
| m=model_outputs["logits"][0].numpy() | |
| e=numpy.exp(m-numpy.max(m,axis=-1,keepdims=True)) | |
| z=e/e.sum(axis=-1,keepdims=True) | |
| for i in range(m.shape[0]-1,0,-1): | |
| m[i-1]+=numpy.nanmax(m[i]+self.transition,axis=1) | |
| k=[numpy.nanargmax(m[0])] | |
| for i in range(1,m.shape[0]): | |
| k.append(numpy.nanargmax(m[i]+self.transition[k[-1]])) | |
| w=[{"entity":self.model.config.id2label[j],"start":s,"end":e,"score":z[i,j]} for i,((s,e),j) in enumerate(zip(model_outputs["offset_mapping"][0].tolist(),k)) if s<e] | |
| if "aggregation_strategy" in kwargs and kwargs["aggregation_strategy"]!="none": | |
| for i,t in reversed(list(enumerate(w))): | |
| p=t.pop("entity") | |
| if p.startswith("I-"): | |
| w[i-1]["score"]=min(w[i-1]["score"],t["score"]) | |
| w[i-1]["end"]=w.pop(i)["end"] | |
| elif p.startswith("B-"): | |
| t["entity_group"]=p[2:] | |
| else: | |
| t["entity_group"]=p | |
| for t in w: | |
| t["text"]=model_outputs["sentence"][t["start"]:t["end"]] | |
| return w | |
| class MistralForTokenClassification(MistralPreTrainedModel): | |
| def __init__(self,config): | |
| from torch import nn | |
| super().__init__(config) | |
| self.num_labels=config.num_labels | |
| self.model=MistralModel(config) | |
| if hasattr(config,"classifier_dropout") and config.classifier_dropout is not None: | |
| classifier_dropout=config.classifier_dropout | |
| elif hasattr(config,"hidden_dropout") and config.hidden_dropout is not None: | |
| classifier_dropout=config.hidden_dropout | |
| else: | |
| classifier_dropout=0.1 | |
| self.dropout=nn.Dropout(classifier_dropout) | |
| self.classifier=nn.Linear(config.hidden_size,config.num_labels) | |
| self.post_init() | |
| def get_input_embeddings(self): | |
| return self.model.embed_tokens | |
| def set_input_embeddings(self,value): | |
| self.model.embed_tokens=value | |
| def forward(self,input_ids=None,past_key_values=None,attention_mask=None,position_ids=None,inputs_embeds=None,labels=None,use_cache=None,output_attentions=None,output_hidden_states=None,return_dict=None): | |
| return_dict=return_dict if return_dict is not None else self.config.use_return_dict | |
| transformer_outputs=self.model(input_ids,past_key_values=past_key_values,attention_mask=attention_mask,position_ids=position_ids,inputs_embeds=inputs_embeds,use_cache=use_cache,output_attentions=output_attentions,output_hidden_states=output_hidden_states,return_dict=return_dict) | |
| hidden_states=transformer_outputs[0] | |
| hidden_states=self.dropout(hidden_states) | |
| logits=self.classifier(hidden_states) | |
| loss=None | |
| if labels is not None: | |
| from torch import nn | |
| loss_fct=nn.CrossEntropyLoss() | |
| loss=loss_fct(logits.view(-1,self.num_labels),labels.view(-1)) | |
| if not return_dict: | |
| output=(logits,)+transformer_outputs[2:] | |
| return ((loss,)+output) if loss is not None else output | |
| return TokenClassifierOutput(loss=loss,logits=logits,hidden_states=transformer_outputs.hidden_states,attentions=transformer_outputs.attentions) | |