File size: 884 Bytes
1d89d18 65266b7 1d89d18 65266b7 1d89d18 65266b7 1d89d18 |
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 |
import gradio as gr
from transformers import AutoModelForTokenClassification, AutoTokenizer
import torch
# Cargar el modelo y el tokenizador
model_name = "EmergentMethods/gliner_medium_news-v2.1"
model = AutoModelForTokenClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def predict(text):
inputs = tokenizer(text, return_tensors="pt")
# Realizar la inferencia
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predictions = torch.argmax(logits, dim=2)
id2label = model.config.id2label
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
entities = [{"token": token, "label": id2label[prediction.item()]} for token, prediction in zip(tokens, predictions[0])]
return entities
demo = gr.Interface(fn=predict, inputs="text", outputs="json")
demo.launch()
|