NuExtract is a version of
phi-3-mini, fine-tuned on a private high-quality synthetic dataset for information extraction.
To use the model, provide an input text (less than 2000 tokens) and a JSON template describing the information you need to extract.
Note: This model is purely extractive, so all text output by the model is present as is in the original text. You can also provide an example of output formatting to help the model understand your task more precisely.
We also provide a tiny(0.5B) and large(7B) version of this model:
NuExtract-tiny and
NuExtract-large
1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4
5def predict_NuExtract(model, tokenizer, text, schema, example=["", "", ""]):
6 schema = json.dumps(json.loads(schema), indent=4)
7 input_llm = "<|input|>\n### Template:\n" + schema + "\n"
8 for i in example:
9 if i != "":
10 input_llm += "### Example:\n"+ json.dumps(json.loads(i), indent=4)+"\n"
11
12 input_llm += "### Text:\n"+text +"\n<|output|>\n"
13 input_ids = tokenizer(input_llm, return_tensors="pt",truncation = True, max_length=4000).to("cuda")
14
15 output = tokenizer.decode(model.generate(**input_ids)[0], skip_special_tokens=True)
16 return output.split("<|output|>")[1].split("<|end-output|>")[0]
17
18
19# We recommend using bf16 as it results in negligable performance loss
20model = AutoModelForCausalLM.from_pretrained("numind/NuExtract", torch_dtype=torch.bfloat16, trust_remote_code=True)
21tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract", trust_remote_code=True)
22
23model.to("cuda")
24
25model.eval()
26
27text = """We introduce Mistral 7B, a 7–billion-parameter language model engineered for
28superior performance and efficiency. Mistral 7B outperforms the best open 13B
29model (Llama 2) across all evaluated benchmarks, and the best released 34B
30model (Llama 1) in reasoning, mathematics, and code generation. Our model
31leverages grouped-query attention (GQA) for faster inference, coupled with sliding
32window attention (SWA) to effectively handle sequences of arbitrary length with a
33reduced inference cost. We also provide a model fine-tuned to follow instructions,
34Mistral 7B – Instruct, that surpasses Llama 2 13B – chat model both on human and
35automated benchmarks. Our models are released under the Apache 2.0 license.
36Code: https://github.com/mistralai/mistral-src
37Webpage: https://mistral.ai/news/announcing-mistral-7b/"""
38
39schema = """{
40 "Model": {
41 "Name": "",
42 "Number of parameters": "",
43 "Number of max token": "",
44 "Architecture": []
45 },
46 "Usage": {
47 "Use case": [],
48 "Licence": ""
49 }
50}"""
51
52prediction = predict_NuExtract(model, tokenizer, text, schema, example=["","",""])
53print(prediction)
54