Update app.py

app.py

@@ -301,20 +301,20 @@ def highlight_words(input_text, json_output):
 
     return highlighted_text
 
-# model = AutoModelForCausalLM.from_pretrained(
-#         "numind/NuExtract-tinyv2",
-#         )
+
 
 model = AutoModelForCausalLM.from_pretrained(
         "numind/NuExtract",
         trust_remote_code=True,
+        device_map="cuda",
+        torch_dtype="auto",
+        attn_implementation="flash_attention_2"
         )
 
 
 tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract")
 tokenizer.eos = tokenizer("<|end-output|>") 
 
-model.to("cuda")
 model.eval()