riccorl's picture
first commit
626eca0
raw
history blame
5.01 kB
import argparse
import torch
from reader.data.relik_reader_sample import load_relik_reader_samples
from relik.reader.pytorch_modules.hf.modeling_relik import (
RelikReaderConfig,
RelikReaderREModel,
)
from relik.reader.relik_reader_re import RelikReaderForTripletExtraction
from relik.reader.utils.relation_matching_eval import StrongMatching
dict_nyt = {
"/people/person/nationality": "nationality",
"/sports/sports_team/location": "sports team location",
"/location/country/administrative_divisions": "administrative divisions",
"/business/company/major_shareholders": "shareholders",
"/people/ethnicity/people": "ethnicity",
"/people/ethnicity/geographic_distribution": "geographic distributi6on",
"/business/company_shareholder/major_shareholder_of": "major shareholder",
"/location/location/contains": "location",
"/business/company/founders": "founders",
"/business/person/company": "company",
"/business/company/advisors": "advisor",
"/people/deceased_person/place_of_death": "place of death",
"/business/company/industry": "industry",
"/people/person/ethnicity": "ethnic background",
"/people/person/place_of_birth": "place of birth",
"/location/administrative_division/country": "country of an administration division",
"/people/person/place_lived": "place lived",
"/sports/sports_team_location/teams": "sports team",
"/people/person/children": "child",
"/people/person/religion": "religion",
"/location/neighborhood/neighborhood_of": "neighborhood",
"/location/country/capital": "capital",
"/business/company/place_founded": "company founded location",
"/people/person/profession": "occupation",
}
def eval(model_path, data_path, is_eval, output_path=None):
if model_path.endswith(".ckpt"):
# if it is a lightning checkpoint we load the model state dict and the tokenizer from the config
model_dict = torch.load(model_path)
additional_special_symbols = model_dict["hyper_parameters"][
"additional_special_symbols"
]
from transformers import AutoTokenizer
from relik.reader.utils.special_symbols import get_special_symbols_re
special_symbols = get_special_symbols_re(additional_special_symbols - 1)
tokenizer = AutoTokenizer.from_pretrained(
model_dict["hyper_parameters"]["transformer_model"],
additional_special_tokens=special_symbols,
add_prefix_space=True,
)
config_model = RelikReaderConfig(
model_dict["hyper_parameters"]["transformer_model"],
len(special_symbols),
training=False,
)
model = RelikReaderREModel(config_model)
model_dict["state_dict"] = {
k.replace("relik_reader_re_model.", ""): v
for k, v in model_dict["state_dict"].items()
}
model.load_state_dict(model_dict["state_dict"], strict=False)
reader = RelikReaderForTripletExtraction(
model, training=False, device="cuda", tokenizer=tokenizer
)
else:
# if it is a huggingface model we load the model directly. Note that it could even be a string from the hub
model = RelikReaderREModel.from_pretrained(model_path)
reader = RelikReaderForTripletExtraction(model, training=False, device="cuda")
samples = list(load_relik_reader_samples(data_path))
for sample in samples:
sample.candidates = [dict_nyt[cand] for cand in sample.candidates]
sample.triplets = [
{
"subject": triplet["subject"],
"relation": {
"name": dict_nyt[triplet["relation"]["name"]],
"type": triplet["relation"]["type"],
},
"object": triplet["object"],
}
for triplet in sample.triplets
]
predicted_samples = reader.read(samples=samples, progress_bar=True)
if is_eval:
strong_matching_metric = StrongMatching()
predicted_samples = list(predicted_samples)
for k, v in strong_matching_metric(predicted_samples).items():
print(f"test_{k}", v)
if output_path is not None:
with open(output_path, "w") as f:
for sample in predicted_samples:
f.write(sample.to_jsons() + "\n")
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_path",
type=str,
default="/home/huguetcabot/alby-re/relik/relik/reader/models/relik_re_reader_base",
)
parser.add_argument(
"--data_path",
type=str,
default="/home/huguetcabot/alby-re/relik/relik/reader/data/testa.jsonl",
)
parser.add_argument("--is-eval", action="store_true")
parser.add_argument("--output_path", type=str, default=None)
args = parser.parse_args()
eval(args.model_path, args.data_path, args.is_eval, args.output_path)
if __name__ == "__main__":
main()