Views
No views yet





1import json
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5def predict_NuExtract(model, tokenizer, texts, template, batch_size=1, max_length=10_000, max_new_tokens=4_000):
6 template = json.dumps(json.loads(template), indent=4)
7 prompts = [f"""<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>""" for text in texts]
8
9 outputs = []
10 with torch.no_grad():
11 for i in range(0, len(prompts), batch_size):
12 batch_prompts = prompts[i:i+batch_size]
13 batch_encodings = tokenizer(batch_prompts, return_tensors="pt", truncation=True, padding=True, max_length=max_length).to(model.device)
14
15 pred_ids = model.generate(**batch_encodings, max_new_tokens=max_new_tokens)
16 outputs += tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
17
18 return [output.split("<|output|>")[1] for output in outputs]
19
20model_name = "numind/NuExtract-v1.5"
21device = "cuda"
22model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device).eval()
23tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
24
25text = """We introduce Mistral 7B, a 7–billion-parameter language model engineered for
26superior performance and efficiency. Mistral 7B outperforms the best open 13B
27model (Llama 2) across all evaluated benchmarks, and the best released 34B
28model (Llama 1) in reasoning, mathematics, and code generation. Our model
29leverages grouped-query attention (GQA) for faster inference, coupled with sliding
30window attention (SWA) to effectively handle sequences of arbitrary length with a
31reduced inference cost. We also provide a model fine-tuned to follow instructions,
32Mistral 7B – Instruct, that surpasses Llama 2 13B – chat model both on human and
33automated benchmarks. Our models are released under the Apache 2.0 license.
34Code: <https://github.com/mistralai/mistral-src>
35Webpage: <https://mistral.ai/news/announcing-mistral-7b/>"""
36
37template = """{
38 "Model": {
39 "Name": "",
40 "Number of parameters": "",
41 "Number of max token": "",
42 "Architecture": []
43 },
44 "Usage": {
45 "Use case": [],
46 "Licence": ""
47 }
48}"""
49
50prediction = predict_NuExtract(model, tokenizer, [text], template)[0]
51print(prediction)
521import json
2
3MAX_INPUT_SIZE = 20_000
4MAX_NEW_TOKENS = 6000
5
6def clean_json_text(text):
7 text = text.strip()
8 text = text.replace("\#", "#").replace("\&", "&")
9 return text
10
11def predict_chunk(text, template, current, model, tokenizer):
12 current = clean_json_text(current)
13
14 input_llm = f"<|input|>\n### Template:\n{template}\n### Current:\n{current}\n### Text:\n{text}\n\n<|output|>" + "{"
15 input_ids = tokenizer(input_llm, return_tensors="pt", truncation=True, max_length=MAX_INPUT_SIZE).to("cuda")
16 output = tokenizer.decode(model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS)[0], skip_special_tokens=True)
17
18 return clean_json_text(output.split("<|output|>")[1])
19
20def split_document(document, window_size, overlap):
21 tokens = tokenizer.tokenize(document)
22 print(f"\tLength of document: {len(tokens)} tokens")
23
24 chunks = []
25 if len(tokens) > window_size:
26 for i in range(0, len(tokens), window_size-overlap):
27 print(f"\t{i} to {i + len(tokens[i:i + window_size])}")
28 chunk = tokenizer.convert_tokens_to_string(tokens[i:i + window_size])
29 chunks.append(chunk)
30
31 if i + len(tokens[i:i + window_size]) >= len(tokens):
32 break
33 else:
34 chunks.append(document)
35 print(f"\tSplit into {len(chunks)} chunks")
36
37 return chunks
38
39def handle_broken_output(pred, prev):
40 try:
41 if all([(v in ["", []]) for v in json.loads(pred).values()]):
42 # if empty json, return previous
43 pred = prev
44 except:
45 # if broken json, return previous
46 pred = prev
47
48 return pred
49
50def sliding_window_prediction(text, template, model, tokenizer, window_size=4000, overlap=128):
51 # split text into chunks of n tokens
52 tokens = tokenizer.tokenize(text)
53 chunks = split_document(text, window_size, overlap)
54
55 # iterate over text chunks
56 prev = template
57 for i, chunk in enumerate(chunks):
58 print(f"Processing chunk {i}...")
59 pred = predict_chunk(chunk, template, prev, model, tokenizer)
60
61 # handle broken output
62 pred = handle_broken_output(pred, prev)
63
64 # iterate
65 prev = pred
66
67 return pred