Update README.md

README.md

@@ -5,7 +5,7 @@ language:
 ---
 # Structure Extraction Model by NuMind 🔥
 
-NuExtract-large is a fine-tuned version of [phi-3-small](https://huggingface.co/microsoft/Phi-3-small-8k-instruct), on a private high-quality syntatic dataset for information extraction. 
+NuExtract-large is a fine-tuned version of [phi-3-small](https://huggingface.co/microsoft/Phi-3-small-8k-instruct), 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 schema 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.
@@ -45,7 +45,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 
-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:
@@ -53,13 +53,13 @@ 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,torch_dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained("numind/NuExtract", trust_remote_code=True, torch_dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract", trust_remote_code=True)
 
 model.to("cuda")
@@ -91,7 +91,7 @@ schema = """{
     }
 }"""
 
-prediction = predict_NuExtract(model,tokenizer,text, schema,example = ["","",""])
+prediction = predict_NuExtract(model, tokenizer, text, schema, example=["","",""])
 print(prediction)
 
 ```