Views
No views yet





1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4def predict_NuExtract(model, tokenizer, texts, template, batch_size=1, max_length=10_000, max_new_tokens=4_000):
5 template = json.dumps(json.loads(template), indent=4)
6 prompts = [f"""<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>""" for text in texts]
7
8 outputs = []
9 with torch.no_grad():
10 for i in range(0, len(prompts), batch_size):
11 batch_prompts = prompts[i:i+batch_size]
12 batch_encodings = tokenizer(batch_prompts, return_tensors="pt", truncation=True, padding=True, max_length=max_length).to(model.device)
13
14 pred_ids = model.generate(**batch_encodings, max_new_tokens=max_new_tokens)
15 outputs += tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
16
17 return [output.split("<|output|>")[1] for output in outputs]
18
19model_name = "numind/NuExtract-v1.5"
20device = "cuda"
21model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device).eval()
22tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
23
24text = """We introduce Mistral 7B, a 7–billion-parameter language model engineered for
25superior performance and efficiency. Mistral 7B outperforms the best open 13B
26model (Llama 2) across all evaluated benchmarks, and the best released 34B
27model (Llama 1) in reasoning, mathematics, and code generation. Our model
28leverages grouped-query attention (GQA) for faster inference, coupled with sliding
29window attention (SWA) to effectively handle sequences of arbitrary length with a
30reduced inference cost. We also provide a model fine-tuned to follow instructions,
31Mistral 7B – Instruct, that surpasses Llama 2 13B – chat model both on human and
32automated benchmarks. Our models are released under the Apache 2.0 license.
33Code: <https://github.com/mistralai/mistral-src>
34Webpage: <https://mistral.ai/news/announcing-mistral-7b/>"""
35
36template = """{
37 "Model": {
38 "Name": "",
39 "Number of parameters": "",
40 "Number of max token": "",
41 "Architecture": []
42 },
43 "Usage": {
44 "Use case": [],
45 "Licence": ""
46 }
47}"""
48
49prediction = predict_NuExtract(model, tokenizer, [text], template)[0]
50print(prediction)
511import 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