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)