Update app.py

app.py

@@ -47,9 +47,42 @@ def get_entities(example):
     
     return output
 
+def clasifica_sistema_universal(example):
+    tokenizer = AutoTokenizer.from_pretrained("hackathon-pln-es/jurisbert-class-tratados-internacionales-sistema-universal")
+
+    model = AutoModelForSequenceClassification.from_pretrained("hackathon-pln-es/jurisbert-class-tratados-internacionales-sistema-universal")
+    text_classifier = pipeline("text-classification",  model=model, tokenizer=tokenizer)
+    results= text_classifier (example)
+
+    salida=[]
+    for i in results:
+      salida.append({i["label"]:i["score"]})
+
+    return {i["label"]: float(i["score"]) for i in results}
+
+def clasifica_conv_americana(example):
+  from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+  tokenizer = AutoTokenizer.from_pretrained("hackathon-pln-es/jurisbert-clas-art-convencion-americana-dh")
+
+  model = AutoModelForSequenceClassification.from_pretrained("hackathon-pln-es/jurisbert-clas-art-convencion-americana-dh")
+  
+  text_classifier = pipeline("text-classification",  model=model, tokenizer=tokenizer)
+  results= text_classifier (example)
+  
+  return {i["label"]: float(i["score"]) for i in results}
+
 def process(example):
     entidades = get_entities(example)
-    return entidades
-  
-iface = gr.Interface(fn=get_entities, inputs="text", outputs=['highlight'], examples=examples, title="Test of jurisbert-finetuning-ner ",)
+
+    class_sistema_universal = clasifica_sistema_universal(example)
+
+    class_conv_americana = clasifica_conv_americana(example)
+
+
+    return entidades,class_sistema_universal, class_conv_americana
+
+
+iface = gr.Interface(fn=process, inputs="text", outputs=["highlight","label","label"], examples=examples, title="Modelo Jurídico Mexicano")
+
 iface.launch(debug=True)