Update app.py

app.py

@@ -5,13 +5,14 @@ import gradio as gr
 import unicodedata
 import torch
 import shutil
+import tempfile
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from langchain_community.document_loaders import PyMuPDFLoader
 
 # Configurações do modelo
 MODEL_PATH = "numind/NuExtract-1.5"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-device = "cuda"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH, 
     torch_dtype=torch.float16  # Usa FP16 para reduzir o uso de VRAM
@@ -68,14 +69,21 @@ def structure_text(text):
     return structured_data
 
 def generate_text(prompt):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # Detecta GPU ou CPU
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_INPUT_SIZE).to(device)  # Move inputs para GPU
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_INPUT_SIZE).to(device)
 
     with torch.no_grad():
-        output = model.generate(**inputs, max_new_tokens=512)
+        output = model.generate(
+            **inputs, 
+            max_new_tokens=512,  # Limita o tamanho da resposta
+            pad_token_id=tokenizer.eos_token_id  # Evita erros na geração
+        )
+
+    torch.cuda.empty_cache()  # Libera VRAM após a geração
 
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
+
+
 def process_chunk(text, template, current):
     input_text = f"### Template:\n{template}\n### Current:\n{current}\n### Text:\n{text}\n"
     output_text = generate_text(input_text)
@@ -84,8 +92,8 @@ def process_chunk(text, template, current):
         parsed_output = json.loads(output_text)
         return json.dumps(parsed_output, indent=2, ensure_ascii=False)
     except json.JSONDecodeError as e:
-        print("Erro ao decodificar JSON:", e)
-        return output_text
+        print(f"[Erro JSON] {e}: {output_text}")  # Log do erro para depuração
+        return json.dumps({"erro": "Saída inválida do modelo", "output_bruto": output_text}, indent=2, ensure_ascii=False)
 
 def handle_broken_outputs(pred, prev):
     try: 
@@ -96,56 +104,50 @@ def handle_broken_outputs(pred, prev):
     return pred
 
 
-
 def send_chunk_to_model(text, template, current):
     """Envia um chunk de texto para o modelo local e processa a saída."""
     input_text = f"<|input|>\n### Template:\n{template}\n### Current:\n{current}\n### Text:\n{text}\n\n<|output|>" + "{"
     output_text = process_chunk(input_text, template, current)
 
-    return handle_broken_output(output_text, current)
+    return handle_broken_outputs(output_text, current)
 
 
 def process_and_generate(pdf_file):
     if not pdf_file:
         return "Nenhum arquivo enviado."
 
-    # Caminho temporário para salvar o arquivo
-    pdf_path = "/tmp/uploaded_file.pdf"
-
-    # Copia o conteúdo do arquivo para o caminho correto
-    shutil.copy(pdf_file, pdf_path)
-
-    # Extrai o texto do PDF
-    extracted_text = extract_text_from_pdf(pdf_path)
-    if not extracted_text:
-        return "Falha ao extrair texto do PDF."
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+    tmp_file.write(pdf_file.read())  # Grava o arquivo diretamente
+    tmp_file.flush()  # Garante que os dados sejam salvos
+    pdf_path = tmp_file.name
 
-    # Estrutura os dados extraídos
-    structured_data = structure_text(extracted_text)
-    
-    # Carrega o template
-    template = json.dumps(load_template(), ensure_ascii=False)
-    current = json.dumps(structured_data, ensure_ascii=False)
+    try:
+        extracted_text = extract_text_from_pdf(pdf_path)
+        if not extracted_text:
+            return "Falha ao extrair texto do PDF."
 
-    # Divide o texto extraído em partes menores
-    chunks = split_document(extracted_text)
+        structured_data = structure_text(extracted_text)
+        template = json.dumps(load_template(), ensure_ascii=False)
+        current = json.dumps(structured_data, ensure_ascii=False)
+        chunks = split_document(extracted_text)
 
-    # Processa cada chunk com o modelo
-    for chunk in chunks:
-        current = send_chunk_to_model(chunk, template, current)
+        for chunk in chunks:
+            current = send_chunk_to_model(chunk, template, current)
 
-    # Retorna o JSON formatado
-    try:
         return json.dumps(json.loads(current), indent=2, ensure_ascii=False)
-    except json.JSONDecodeError:
-        return "Erro ao processar os dados gerados pelo modelo."
+    except Exception as e:
+        return f"Erro durante o processamento: {e}"
+    finally:
+        os.remove(pdf_path)  # Remove o arquivo temporário após o uso
+
 
 
 interface = gr.Interface(
     fn=process_and_generate,
     inputs=gr.File(label="Upload PDF"),
-    outputs="text",
+    outputs=gr.JSON(label="Dados Extraídos"),
     title="Extração de Dados com Modelo Local",
-    description="Extrai texto de PDFs e processa utilizando um modelo local de linguagem."
+    description="Envie um PDF para extrair e processar informações automaticamente.",
 )
+
 interface.launch()