File size: 3,915 Bytes
7a3c511 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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 RawTokenClassificationPipeline(TokenClassificationPipeline):
def check_model_type(self,supported_models):
pass
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)
|