Views
No views yet



1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4def triplextract(model, tokenizer, text, entity_types, predicates):
5
6 input_format = """
7 **Entity Types:**
8 {entity_types}
9
10 **Predicates:**
11 {predicates}
12
13 **Text:**
14 {text}
15 """
16
17 message = input_format.format(
18 entity_types = json.dumps({"entity_types": entity_types}),
19 predicates = json.dumps({"predicates": predicates}),
20 text = text)
21
22 messages = [{'role': 'user', 'content': message}]
23 input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt = True, return_tensors="pt").to("cuda")
24 output = tokenizer.decode(model.generate(input_ids=input_ids, max_length=2048)[0], skip_special_tokens=True)
25 return output
26
27model = AutoModelForCausalLM.from_pretrained("sciphi/triplex", trust_remote_code=True).to('cuda').eval()
28tokenizer = AutoTokenizer.from_pretrained("sciphi/triplex", trust_remote_code=True)
29
30entity_types = [ "LOCATION", "POSITION", "DATE", "CITY", "COUNTRY", "NUMBER" ]
31predicates = [ "POPULATION", "AREA" ]
32text = """
33San Francisco,[24] officially the City and County of San Francisco, is a commercial, financial, and cultural center in Northern California.
34
35With 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.
36"""
37
38prediction = triplextract(model, tokenizer, text, entity_types, predicates)
39print(prediction)
40
41