app.py

import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
# Usaremos AutoTokenizer e AutoModelForCausalLM para o novo modelo
from transformers import AutoTokenizer, AutoModelForCausalLM
from pypdf import PdfReader
import os
import re # MÓDULO ADICIONADO para Expressões Regulares

# --- 1. Carregamento dos Modelos ---
# Modelo de recuperação não muda, ele é excelente para essa tarefa.
print("Carregando o modelo de recuperação (Sentence Transformer)...")
retriever_model = SentenceTransformer('all-MiniLM-L6-v2')

# Carregando o modelo de geração DeepSeek
print("Carregando o modelo de geração (DeepSeek)...")
# Nota: "trust_remote_code=True" é necessário para carregar a arquitetura do DeepSeek
generator_tokenizer = AutoTokenizer.from_pretrained(
    'deepseek-ai/deepseek-coder-1.3b-instruct', 
    trust_remote_code=True
)
generator_model = AutoModelForCausalLM.from_pretrained(
    'deepseek-ai/deepseek-coder-1.3b-instruct', 
    trust_remote_code=True
)
print("Modelos carregados com sucesso!")

# --- 2. Função para Processar Arquivos Enviados (COM CHUNKING ESTRUTURADO) ---
def process_files(files):
    if not files:
        return None, "Por favor, envie um ou mais arquivos."
    knowledge_text = ""
    for file in files:
        file_path = file.name
        if file_path.endswith(".pdf"):
            try:
                reader = PdfReader(file_path)
                for page in reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        # Adiciona um espaço extra entre as páginas para garantir a separação
                        knowledge_text += page_text + "\n\n"
            except Exception as e:
                return None, f"Erro ao ler o arquivo PDF {os.path.basename(file_path)}: {e}"
        elif file_path.endswith(".txt"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    knowledge_text += f.read() + "\n\n"
            except Exception as e:
                return None, f"Erro ao ler o arquivo TXT {os.path.basename(file_path)}: {e}"

    if not knowledge_text.strip():
        return None, "Não foi possível extrair texto dos arquivos fornecidos."

    # MUDANÇA PRINCIPAL: Extrator Estruturado usando Regex
    # Este padrão de regex divide o texto ANTES de uma linha que parece um cabeçalho de seção
    # (Ex: "1. CLIENTE", "8. PADRÕES"). Isso mantém a seção inteira em um único chunk.
    # O `(?m)` ativa o modo multiline, fazendo `^` corresponder ao início de cada linha.
    chunk_pattern = r"(?m)(^\d+\..*)"
    
    # Divide o texto em chunks usando o padrão e remove os vazios
    text_chunks = [chunk.strip() for chunk in re.split(chunk_pattern, knowledge_text) if chunk.strip()]
    
    # Reagrupa o cabeçalho com seu conteúdo
    structured_chunks = []
    i = 0
    while i < len(text_chunks):
        if re.match(chunk_pattern, text_chunks[i]) and i + 1 < len(text_chunks):
            # Junta o cabeçalho (ex: "1. CLIENTE") com o conteúdo seguinte
            structured_chunks.append(text_chunks[i] + "\n" + text_chunks[i+1])
            i += 2
        else:
            # Adiciona conteúdo que não corresponde a um cabeçalho
            structured_chunks.append(text_chunks[i])
            i += 1

    if not structured_chunks:
        return None, "O texto extraído não continha blocos de texto válidos para processamento."

    print(f"Processando {len(structured_chunks)} chunks estruturados dos arquivos...")
    knowledge_base_embeddings = retriever_model.encode(structured_chunks, convert_to_tensor=True, show_progress_bar=True)
    print("Base de conhecimento criada a partir dos arquivos.")
    
    return (structured_chunks, knowledge_base_embeddings), f"✅ Sucesso! {len(files)} arquivo(s) processado(s), gerando {len(structured_chunks)} chunks estruturados."


# --- 3. A Função Principal do RAG (sem alterações) ---
def answer_question(question, knowledge_state):
    if not question:
        return "Por favor, insira uma pergunta."
    if not knowledge_state or not knowledge_state[0] or knowledge_state[1] is None:
        return "⚠️ A base de conhecimento está vazia. Por favor, processe alguns arquivos primeiro."

    knowledge_base, knowledge_base_embeddings = knowledge_state

    # Etapa de Recuperação
    question_embedding = retriever_model.encode(question, convert_to_tensor=True)
    cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
    # Aumentado para 7 para mais contexto
    top_k = min(7, len(knowledge_base)) 
    top_results = torch.topk(cosine_scores, k=top_k, dim=-1)
    retrieved_context = "\n---\n".join([knowledge_base[i] for i in top_results.indices[0]])

    if not retrieved_context.strip():
        return "Não foi possível encontrar um contexto relevante nos documentos para responder a esta pergunta."

    print(f"\n--- Nova Pergunta de Auditoria ---")
    print(f"Pergunta: {question}")
    print(f"Contexto Recuperado (Top {top_k}):\n{retrieved_context}")

    # Prompt com regras explícitas de extração de entidades
    prompt = f"""### Instruction:
Você é um assistente de IA especialista em extrair informações de documentos técnicos. Analise o 'Contexto' para responder à 'Pergunta' seguindo estas regras rigorosamente:

**Regras de Extração:**
1.  **Use APENAS a informação do 'Contexto'.** Não adicione informações externas.
2.  **Para perguntas sobre 'cliente':** Procure por linhas que comecem com "Cliente:", "Contratante:", ou um nome de empresa claro. Ignore frases genéricas sobre escopo ou lotes.
3.  **Para perguntas sobre 'instrumento':** Procure por linhas que comecem com "Instrumento:", "Descrição:", "Objeto:", "TAG:", ou um modelo específico. Ignore descrições de métodos de calibração.
4.  **Para perguntas sobre 'título' ou 'resumo':** Resuma os dados principais, como o tipo de documento (ex: Certificado de Calibração), o nome do cliente e o instrumento calibrado.

**Contexto:**
{retrieved_context}

**Pergunta:**
{question}

### Response:
"""

    input_ids = generator_tokenizer(prompt, return_tensors="pt").input_ids
    input_length = input_ids.shape[1]

    # Ajuste nos parâmetros de geração
    outputs = generator_model.generate(
        input_ids,
        # Aumentado para permitir respostas mais detalhadas
        max_new_tokens=350,  
        do_sample=False,
        eos_token_id=generator_tokenizer.eos_token_id,
        pad_token_id=generator_tokenizer.eos_token_id
    )

    # Decodificação correta para modelos Causal LM
    generated_tokens = outputs[0, input_length:]
    answer = generator_tokenizer.decode(generated_tokens, skip_special_tokens=True)
    
    return answer

# --- 4. Interface Gráfica (sem alterações na estrutura) ---
with gr.Blocks(theme=gr.themes.Soft()) as interface:
    knowledge_state = gr.State()
    gr.Markdown(
        """
        # 🤖 RAG - Auditor de Documentos (v9 - Chunking Estruturado)
        **1. Carregue seus arquivos**: Envie um ou mais certificados ou documentos nos formatos `.pdf` ou `.txt`.
        **2. Processe os arquivos**: Clique no botão para criar a base de conhecimento.
        **3. Faça perguntas**: Após o processamento, faça perguntas sobre o conteúdo dos documentos.
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            file_uploader = gr.File(label="Carregar Certificados (.pdf, .txt)", file_count="multiple", file_types=[".pdf", ".txt"])
            process_button = gr.Button("Processar Arquivos", variant="primary")
            status_box = gr.Textbox(label="Status do Processamento", interactive=False)
        with gr.Column(scale=2):
            question_box = gr.Textbox(label="Faça sua pergunta aqui", placeholder="Ex: Qual o resultado da calibração do instrumento PI-101?")
            submit_button = gr.Button("Obter Resposta", variant="primary")
            answer_box = gr.Textbox(label="Resposta Baseada nos Documentos", interactive=False, lines=5)

    process_button.click(fn=process_files, inputs=[file_uploader], outputs=[knowledge_state, status_box])
    submit_button.click(fn=answer_question, inputs=[question_box, knowledge_state], outputs=[answer_box])

# --- 5. Lançamento do App ---
if __name__ == "__main__":
    interface.launch()