thiagoloth/force-ia

1import os2import glob3import gc4import gradio as gr5from typing import List6 7# -------- LangChain + embeddings + FAISS --------8from langchain_community.document_loaders import PyPDFLoader9from langchain_text_splitters import RecursiveCharacterTextSplitter10from langchain_community.vectorstores import FAISS11from langchain_community.embeddings import HuggingFaceEmbeddings12 13# -------- Modelo gerativo --------14from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline15 16# -------- STT: voz -> texto --------17from faster_whisper import WhisperModel18 19# -------- Configurações --------20PDF_DIR = "PDFs"21INDEX_PATH = "faiss_index"22EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"23GEN_MODEL_NAME = "google/flan-t5-small"  # versão leve24STT_MODEL_SIZE = "tiny"  # Whisper tiny para CPU25 26# -------- Funções auxiliares --------27def list_pdf_files() -> List[str]:28    os.makedirs(PDF_DIR, exist_ok=True)29    return sorted(glob.glob(os.path.join(PDF_DIR, "*.pdf")))30 31def load_documents(pdf_paths: List[str]):32    docs = []33    for path in pdf_paths:34        try:35            loader = PyPDFLoader(path)36            docs.extend(loader.load())37        except Exception as e:38            print(f"[WARN] Falha ao carregar {path}: {e}")39    return docs40 41def build_or_load_index(force_rebuild: bool = False):42    pdfs = list_pdf_files()43    if not pdfs:44        return None, [], "Nenhum PDF encontrado na pasta PDFs/."45 46    if (not force_rebuild) and os.path.exists(INDEX_PATH):47        try:48            embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)49            vs = FAISS.load_local(INDEX_PATH, embeddings, allow_dangerous_deserialization=True)50            return vs, pdfs, f"Index carregado. PDFs indexados: {len(pdfs)}."51        except Exception as e:52            print(f"[INFO] Não foi possível carregar index existente: {e}. Recriando...")53 54    # Build do zero55    docs = load_documents(pdfs)56    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)57    chunks = splitter.split_documents(docs)58 59    embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)60    vs = FAISS.from_documents(chunks, embeddings)61 62    os.makedirs(INDEX_PATH, exist_ok=True)63    vs.save_local(INDEX_PATH)64    return vs, pdfs, f"Index criado com {len(chunks)} chunks a partir de {len(pdfs)} PDFs."65 66# -------- Modelo gerativo --------67_tokenizer = None68_model = None69_gen = None70def load_generator():71    global _tokenizer, _model, _gen72    if _gen is not None:73        return _gen74    _tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL_NAME)75    _model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL_NAME)76    _gen = pipeline(77        "text2text-generation",78        model=_model,79        tokenizer=_tokenizer,80        max_new_tokens=384,81        do_sample=False,82        num_beams=4,83    )84    return _gen85 86# -------- STT --------87_whisper = None88def load_stt():89    global _whisper90    if _whisper is not None:91        return _whisper92    _whisper = WhisperModel(STT_MODEL_SIZE, device="cpu", compute_type="int8")93    return _whisper94 95def transcribe(audio_path: str) -> str:96    stt = load_stt()97    if audio_path is None or not os.path.exists(audio_path):98        return ""99    segments, _ = stt.transcribe(audio_path, language="pt")100    text_parts = [seg.text.strip() for seg in segments]101    return " ".join(text_parts).strip()102 103# -------- RAG --------104SYSTEM_INSTRUCTIONS = (105    "Você é um assistente experto em empreendimentos imobiliários. Responda SEMPRE em português do Brasil, "106    "com tom claro e profissional. Use o contexto dos PDFs como fonte principal. "107    "Você PODE criar textos gerativos (pitches, resumos, copy de vendas), mas deve manter coerência com os dados."108)109 110PROMPT_TEMPLATE = (111    "{instrucoes}\n\nPergunta: {pergunta}\n\nContexto (trechos relevantes dos PDFs):\n{contexto}\n\nResposta:"112)113 114def retrieve_context(vs: FAISS, query: str, k: int = 6) -> str:115    retriever = vs.as_retriever(search_type="similarity", search_kwargs={"k": k})116    docs = retriever.get_relevant_documents(query)117    return "\n\n".join([d.page_content for d in docs])  # somente texto limpo118 119def generate_answer(vs: FAISS, query: str) -> str:120    gen = load_generator()121    contexto = retrieve_context(vs, query)122    prompt = PROMPT_TEMPLATE.format(instrucoes=SYSTEM_INSTRUCTIONS, pergunta=query, contexto=contexto)123    out = gen(prompt)[0]["generated_text"].strip()124    return out.replace("\n\n\n", "\n\n").strip()125 126# -------- Gradio UI --------127with gr.Blocks(css=".gradio-container {max-width: 900px !important}") as demo:128    gr.Markdown("""129    # 🏡 Chat de Imóveis por Voz ou Texto130    **Pergunte por voz ou digite, receba a resposta em texto.**131    Suba seus PDFs em `PDFs/` e clique em 🔁 Reindexar PDFs.132    """)133 134    vs_state = gr.State()135    status = gr.Markdown("Inicializando…")136 137    # Entradas138    audio = gr.Audio(sources=["microphone"], type="filepath", label="Pergunte por voz (pt-BR)")139    text_input = gr.Textbox(label="Ou digite sua pergunta", placeholder="Digite aqui sua pergunta...", lines=1)140    chat = gr.Chatbot(height=360)141 142    # Botões143    rebuild_btn = gr.Button("🔁 Reindexar PDFs")144    clear_btn = gr.Button("🧹 Limpar chat")145 146    # Inicialização e funções147    def _init():148        vs, pdfs, msg = build_or_load_index(force_rebuild=False)149        load_generator()150        return vs, f"**Status do índice:** {msg}\nPDFs detectados: {len(pdfs)}"151 152    def _rebuild():153        vs, pdfs, msg = build_or_load_index(force_rebuild=True)154        return vs, f"**Status do índice:** {msg}\nPDFs detectados: {len(pdfs)}"155 156    def _ask_audio(audio_path, history, vs):157        if vs is None:158            return history + [["(voz)", "⚠️ Nenhum índice carregado. Adicione PDFs e clique em Reindexar."]]159        query = transcribe(audio_path)160        if not query:161            return history + [["(voz)", "Não consegui ouvir. Tente novamente."]]162        ans = generate_answer(vs, query)163        history = history + [[query, ans]]164        gc.collect()165        return history166 167    def _ask_text(query, history, vs):168        if vs is None:169            return history + [["(texto)", "⚠️ Nenhum índice carregado. Adicione PDFs e clique em Reindexar."]]170        if not query.strip():171            return history + [["(texto)", "Digite uma pergunta válida."]]172        ans = generate_answer(vs, query)173        history = history + [[query, ans]]174        gc.collect()175        return history176 177    # Inicializa178    vs_init, stat_msg = _init()179    vs_state.value = vs_init180    status.value = stat_msg181 182    # Eventos183    rebuild_btn.click(_rebuild, inputs=None, outputs=[vs_state, status])184    clear_btn.click(lambda: [], inputs=None, outputs=chat)185    audio.change(_ask_audio, inputs=[audio, chat, vs_state], outputs=chat)186    text_input.submit(_ask_text, inputs=[text_input, chat, vs_state], outputs=chat)187 188if __name__ == "__main__":189    demo.launch()190