Views
No views yet
1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4def triplextract(model, tokenizer, text, entity_types, predicates):
5
6 input_format = """Perform Named Entity Recognition (NER) and extract knowledge graph triplets from the text. NER identifies named entities of given entity types, and triple extraction identifies relationships between entities using specified predicates.
7
8 **Entity Types:**
9 {entity_types}
10
11 **Predicates:**
12 {predicates}
13
14 **Text:**
15 {text}
16 """
17
18 message = input_format.format(
19 entity_types = json.dumps({"entity_types": entity_types}),
20 predicates = json.dumps({"predicates": predicates}),
21 text = text)
22
23 messages = [{'role': 'user', 'content': message}]
24 input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt = True, return_tensors="pt").to("cuda")
25 output = tokenizer.decode(model.generate(input_ids=input_ids, max_length=2048)[0], skip_special_tokens=True)
26 return output
27
28model = AutoModelForCausalLM.from_pretrained("sciphi/triplex", trust_remote_code=True).to('cuda').eval()
29tokenizer = AutoTokenizer.from_pretrained("sciphi/triplex", trust_remote_code=True)
30
31entity_types = [ "LOCATION", "POSITION", "DATE", "CITY", "COUNTRY", "NUMBER" ]
32predicates = [ "POPULATION", "AREA" ]
33text = """
34San Francisco,[24] officially the City and County of San Francisco, is a commercial, financial, and cultural center in Northern California.
35
36With a population of 808,437 residents as of 2022, San Francisco is the fourth most populous city in the U.S. state of California behind Los Angeles, San Diego, and San Jose.
37"""
38
39prediction = triplextract(model, tokenizer, text, entity_types, predicates)
40print(prediction)
41
42@misc{pimpalgaonkar2024triplex,
author = {Pimpalgaonkar, Shreyas and Tremelling, Nolan and Colegrove, Owen},
title = {Triplex: a SOTA LLM for knowledge graph construction},
year = {2024},
url = {https://huggingface.co/sciphi/triplex}
}