Update README.md

README.md

@@ -5,12 +5,12 @@ language:
 ---
 # Structure Extraction Model by NuMind 🔥
 
-NuExtract is a fine-tuned version of [phi-3-mini](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct), on a private high-quality syntatic dataset for information extraction. 
-To use the model, provide an input text (less than 2000 tokens) and a JSON schema describing the information you need to extract. 
+NuExtract is a version of [phi-3-mini](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct), 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 each information output by the model is present as it is in the text. You can also provide an example of output to help the model understand your task more precisely.
+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.
 
-try here: https://huggingface.co/spaces/numind/NuExtract
+Try it here: https://huggingface.co/spaces/numind/NuExtract
 
 We also provide a tiny(0.5B) and large(7B) version of this model: [NuExtract-tiny](https://huggingface.co/numind/NuExtract-tiny) and [NuExtract-large](https://huggingface.co/numind/NuExtract-large)
 
@@ -44,7 +44,7 @@ import json
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
-def predict_NuExtract(model,tokenizer,text, schema,example = ["","",""]):
+def predict_NuExtract(model, tokenizer, text, schema, example=["", "", ""]):
     schema = json.dumps(json.loads(schema), indent=4)
     input_llm =  "<|input|>\n### Template:\n" +  schema + "\n"
     for i in example:
@@ -52,13 +52,14 @@ def predict_NuExtract(model,tokenizer,text, schema,example = ["","",""]):
           input_llm += "### Example:\n"+ json.dumps(json.loads(i), indent=4)+"\n"
     
     input_llm +=  "### Text:\n"+text +"\n<|output|>\n"
-    input_ids = tokenizer(input_llm, return_tensors="pt",truncation = True, max_length = 4000).to("cuda")
+    input_ids = tokenizer(input_llm, return_tensors="pt",truncation = True, max_length=4000).to("cuda")
 
     output = tokenizer.decode(model.generate(**input_ids)[0], skip_special_tokens=True)
     return output.split("<|output|>")[1].split("<|end-output|>")[0]
 
 
-model = AutoModelForCausalLM.from_pretrained("numind/NuExtract", trust_remote_code=True)
+# We recommend using bf16 as it results in negligable performance loss
+model = AutoModelForCausalLM.from_pretrained("numind/NuExtract", torch_dtype=torch.bfloat16, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract", trust_remote_code=True)
 
 model.to("cuda")
@@ -90,7 +91,7 @@ schema = """{
     }
 }"""
 
-prediction = predict_NuExtract(model,tokenizer,text, schema,example = ["","",""])
+prediction = predict_NuExtract(model, tokenizer, text, schema, example=["","",""])
 print(prediction)
 
 ```