Update README.md

README.md

@@ -43,6 +43,7 @@ To use the model:
 
 ```python
 import json
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 def predict_NuExtract(model, tokenizer, texts, template, batch_size=1, max_length=10_000, max_new_tokens=4_000):
@@ -60,7 +61,7 @@ def predict_NuExtract(model, tokenizer, texts, template, batch_size=1, max_lengt
 
     return [output.split("<|output|>")[1] for output in outputs]
 
-model_name = "numind/NuExtract-v1.5"
+model_name = "numind/NuExtract-1.5-smol"
 device = "cuda"
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, trust_remote_code=True).to(device).eval()
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)