File size: 1,633 Bytes
178c5c2 5bc7c1a e8e65ad 5bc7c1a 0e95eb5 178c5c2 678fb1e fe4bf72 678fb1e fe4bf72 178c5c2 678fb1e 5bc7c1a fe4bf72 478d286 5bc7c1a |
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 |
import gradio as gr
import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
peft_model_id = "hackathon-somos-nlp-2023/bertin-gpt-j-6b-ner-es"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(
config.base_model_name_or_path,
return_dict=True,
load_in_8bit=True,
device_map="auto",
revision="half",
)
tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)
model.eval()
def gen_entities(in_text):
"""Does Named Entity Recognition in the given text."""
text = f"<SP> text: {in_text}\n\n entities:"
batch = tokenizer(text, return_tensors="pt")
batch["input_ids"] = batch["input_ids"].to("cuda")
with torch.cuda.amp.autocast():
output_tokens = model.generate(**batch, max_new_tokens=256, eos_token_id=50258)
response = tokenizer.batch_decode(
output_tokens.detach().cpu().numpy(), skip_special_tokens=False
)[0]
return response[response.find("entities") : response.find("<EP>")]
iface = gr.Interface(
fn=gen_entities,
inputs="text",
outputs="text",
title="Podcast Named Entity Recognition",
description="Introduce un texto corto para que el modelo identifique las identidades presentes en el mismo.",
theme="gradio/monochrome",
examples=[
[
"Yo hoy voy a hablar de mujeres en el mundo del arte, porque me ha leído un libro fantástico que se llama Historia del arte sin hombres, de Katie Hesel."
]
],
)
iface.launch()
|