Views
No views yet


1import json
2from transformers import AutoModelForCausalLM, AutoTokenizer
3import torch
4
5
6def predict_NuExtract(model, tokenizer, text, schema, example=["","",""]):
7 schema = json.dumps(json.loads(schema), indent=4)
8 input_llm = "<|input|>\n### Template:\n" + schema + "\n"
9 for i in example:
10 if i != "":
11 input_llm += "### Example:\n"+ json.dumps(json.loads(i), indent=4)+"\n"
12
13 input_llm += "### Text:\n"+text +"\n<|output|>\n"
14 input_ids = tokenizer(input_llm, return_tensors="pt", truncation=True, max_length=4000).to("cuda")
15
16 output = tokenizer.decode(model.generate(**input_ids)[0], skip_special_tokens=True)
17 return output.split("<|output|>")[1].split("<|end-output|>")[0]
18
19
20model = AutoModelForCausalLM.from_pretrained("numind/NuExtract", trust_remote_code=True, torch_dtype=torch.bfloat16)
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 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