This model is a sub-task of text-to-json task that generates a JSON template given a text.
1import json
2import torch
3from transformers import AutoModel, AutoTokenizer
4
5model_name = "chnaaam/luSI-v1.0"
6
7if torch.cuda.is_available():
8 device = "cuda"
9elif torch.backends.mps.is_available():
10 device = "mps"
11else:
12 device = "cpu"
13
14model = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device).eval()
15tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
16
17text = """아이유(IU, 본명: 이지은, 李知恩[1], 1993년 5월 16일~)는 대한민국의 싱어송라이터, 작곡가, 배우이다. 2007년 로엔 엔터테인먼트(현 카카오 엔터테인먼트) 연습생으로 전속 계약을 맺고 15세의 나이에 2008년 첫 EP인 로스트 앤 파운드(Lost and Found)를 통해 가수로 데뷔했다."""
18
19messages = [
20 {"role": "user", "content": text}
21]
22text = tokenizer.apply_chat_template(
23 messages,
24 tokenize=False,
25 add_generation_prompt=True
26)
27
28model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
29
30generated_ids = model.generate(
31 **model_inputs,
32 max_new_tokens=1024,
33 temperature=0.0
34)
35generated_ids = [
36 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
37]
38
39response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
40
41json_template = json.loads(response)
42
43print(json_template)
1{
2 'Person': {
3 'Name': '',
4 'Stage name': '',
5 'Real name': '',
6 'Birth date': '',
7 'Nationality': '',
8 'Occupations': [],
9 'Debut': {
10 'Age': '',
11 'Year': '',
12 'Company': '',
13 'Contract type': '',
14 'EP': '',
15 'EP title': ''
16 }
17 }
18}