gyr66

/

relation_extraction_bert_base_uncased

like 4

what I have atm
I'm simply unsure how to setup the prompt and extract the proper label(s).

My guess atm is to use a mask token after a sep at the end?

from transformers import BertTokenizer, BertForTokenClassification
import torch
import torch.nn.functional as F

# Load the tokenizer and model
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForTokenClassification.from_pretrained('gyr66/relation_extraction_bert_base_uncased',ignore_mismatched_sizes=True)

# Example text, make sure the entities are correctly marked if that's required by the model
text = "The system as described above has its greatest application in an arrayed <e1>configuration</e1> of antenna <e2>elements</e2>."

# Encode the text using the tokenizer
inputs = tokenizer(text, return_tensors="pt")

# Perform inference
model.eval()  # Set the model to evaluation mode
with torch.no_grad():  # Turn off gradients for prediction, saves memory and computations
    outputs = model(**inputs)
    logits = outputs.logits  # Access logits

# Decode predictions
predictions = torch.argmax(logits, dim=-1)

# Output predicted relation tags for each token
predicted_tags = [model.config.id2label[prediction.item()] for prediction in predictions[0]]

# Compute the probabilities for each tag
probs = F.softmax(logits, dim=-1)

tag_probs = {tag: 0 for tag in set(predicted_tags)}
for i, tag in enumerate(predicted_tags):
    tag_probs[tag] += probs[0][i][model.config.label2id[tag]]

for tag, prob in tag_probs.items():
    print(f"{tag}: {prob}")