app.py

import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
# MUDANÇA: Usaremos AutoTokenizer e AutoModelForCausalLM para o novo modelo
from transformers import AutoTokenizer, AutoModelForCausalLM
from pypdf import PdfReader
import os

# --- 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')

# MUDANÇA: Carregando o modelo 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 (sem alterações) ---
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:
                        knowledge_text += page_text + "\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"
            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."

    text_chunks = [chunk.strip() for chunk in knowledge_text.split('\n') if chunk.strip() and len(chunk) > 10]
    if not text_chunks:
        return None, "O texto extraído não continha blocos de texto válidos para processamento."

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

# --- 3. A Função Principal do RAG (com prompt e decodificação ajustados) ---
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 (sem alterações)
    question_embedding = retriever_model.encode(question, convert_to_tensor=True)
    cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
    top_k = min(3, 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]])

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

    # MUDANÇA: Prompt ajustado para o formato de instrução do DeepSeek
    prompt = f"""### Instruction:
Você é um assistente de auditoria especialista. Sua tarefa é sintetizar as informações dos documentos fornecidos para responder à pergunta do auditor. Elabore uma resposta completa e concisa.

**Documentos:**
{retrieved_context}

**Pergunta:**
{question}

### Response:
"""

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

    # MUDANÇA: Ajuste nos parâmetros de geração
    outputs = generator_model.generate(
        input_ids,
        max_new_tokens=256,  # Controla o tamanho da *nova* resposta gerada
        do_sample=False,     # Desativa a amostragem para respostas mais diretas
        eos_token_id=generator_tokenizer.eos_token_id,
        pad_token_id=generator_tokenizer.eos_token_id # Evita warnings
    )

    # MUDANÇA: Decodificação correta para modelos Causal LM
    # Precisamos remover o prompt inicial da saída gerada.
    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 (v5 - DeepSeek)
        **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()