add cuda

app.py

@@ -9,12 +9,16 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # MODEL_NAME = "numind/NuExtract-1.5-tiny"
 MODEL_NAME = "numind/NuExtract-1.5"
-DEVICE = "cpu"
+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)
 
-MODEL = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, trust_remote_code=True)
+MODEL = AutoModelForCausalLM.from_pretrained(MODEL_NAME, 
+                                             torch_dtype=torch.bfloat16, 
+                                             trust_remote_code=True,
+                                             device_map="auto")
 TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+MODEL.eval()
 
 MAX_INPUT_SIZE = 20_000
 MAX_NEW_TOKENS = 6000
@@ -64,11 +68,6 @@ def process_and_generate(pdf_file):
     print(template)
     print(current)
     
-    # chunks = split_document(extracted_text)
-    # for chunk in chunks:
-    #     current = send_chunk_to_model(chunk, template, current)
-    # return json.dumps(json.loads(current), indent=2, ensure_ascii=False)
-    
     pred_template = sliding_window_prediction(current, template, MODEL, TOKENIZER)
     return pred_template
 
@@ -125,4 +124,4 @@ interface = gr.Interface(
     description="Extrai informações do PDF, preenche um modelo JSON e gera perguntas sobre o conteúdo usando Mistral."
 )
 
-interface.launch(share=True)
+interface.launch(debug=True, share=True)