Update app.py

app.py

@@ -1,179 +1,165 @@
 import gradio as gr
 import os
+import torch
 
 from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.llms import HuggingFacePipeline
+from langchain_huggingface import HuggingFacePipeline
 from langchain.memory import ConversationBufferMemory
 
 from transformers import AutoTokenizer, pipeline
-import torch
 
-# Lista de modelos 100% abertos e gratuitos
-list_llm = [
-    "google/flan-t5-xxl",
-    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-    "microsoft/phi-2",
-    "facebook/opt-1.3b",
-    "EleutherAI/gpt-neo-1.3B",
-    "bigscience/bloom-1b7",
-    "RWKV/rwkv-4-169m-pile",
-    "gpt2-medium",
-    "databricks/dolly-v2-3b",
-    "mosaicml/mpt-7b-instruct"
-]
+# ===================================================================
+# CONFIGURAÇÃO RADICAL DE HARDWARE
+# ===================================================================
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TORCH_DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
+MAX_MEMORY = "16GB" if DEVICE == "cpu" else None
 
-list_llm_simple = [name.split("/")[-1] for name in list_llm]
+# ===================================================================
+# LISTA DE MODELOS OTIMIZADOS
+# ===================================================================
+LLM_MODELS = {
+    "TinyLlama-1.1B-Chat": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "Phi-2": "microsoft/phi-2",
+    "Mistral-7B-Instruct": "mistralai/Mistral-7B-Instruct-v0.2",
+    "Zephyr-7B-Beta": "HuggingFaceH4/zephyr-7b-beta"
+}
 
-# Função para carregar documentos PDF
-def load_doc(list_file_path, chunk_size, chunk_overlap):
-    loaders = [PyPDFLoader(file_path) for file_path in list_file_path]
-    pages = []
-    for loader in loaders:
-        pages.extend(loader.load())
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=chunk_size,
-        chunk_overlap=chunk_overlap
-    )
-    return text_splitter.split_documents(pages)
+# ===================================================================
+# NÚCLEO DO SISTEMA
+# ===================================================================
+class DocumentProcessor:
+    @staticmethod
+    def load_and_split(files, chunk_size=512, chunk_overlap=64):
+        """Carrega e processa documentos com paralelismo extremo"""
+        try:
+            loaders = [PyPDFLoader(file.name) for file in files]
+            return [page for loader in loaders for page in loader.load_and_split(
+                RecursiveCharacterTextSplitter(
+                    chunk_size=chunk_size,
+                    chunk_overlap=chunk_overlap,
+                    separators=["\n\n", "\n", " ", ""]
+                )
+            )]
+        except Exception as e:
+            raise RuntimeError(f"FALHA CRÍTICA NO PROCESSAMENTO: {str(e)}")
 
-# Função para criar banco de dados vetorial
-def create_db(splits, collection_name):
-    embedding = HuggingFaceEmbeddings()
-    return Chroma.from_documents(
-        documents=splits,
-        embedding=embedding,
-        persist_directory=f"./{collection_name}"
-    )
+class VectorDBManager:
+    @staticmethod
+    def create(splits):
+        """Cria vetorização com aceleração de hardware"""
+        return Chroma.from_documents(
+            documents=splits,
+            embedding=HuggingFaceEmbeddings(),
+            persist_directory="./chroma_db"
+        )
 
-# Função para inicializar o modelo LLM
-def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
-    progress(0.1, desc="Carregando tokenizer...")
-    
-    tokenizer = AutoTokenizer.from_pretrained(llm_model)
-    
-    progress(0.4, desc="Inicializando pipeline...")
-    
-    # Define a tarefa correta para cada modelo
-    task = "text2text-generation" if "flan-t5" in llm_model.lower() else "text-generation"
-    
-    # Configuração específica para dispositivos
-    device = 0 if torch.cuda.is_available() else -1
-    if "phi-2" in llm_model.lower() and device == 0:
-        device = "cuda"
-    
-    pipeline_obj = pipeline(
-        task,
-        model=llm_model,
-        tokenizer=tokenizer,
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        device=device,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        top_k=top_k,
-        temperature=temperature
-    )
-    
-    llm = HuggingFacePipeline(pipeline=pipeline_obj)
-    
-    progress(0.7, desc="Configurando memória...")
-    memory = ConversationBufferMemory(
-        memory_key="chat_history",
-        return_messages=True
-    )
-    
-    progress(0.8, desc="Criando cadeia...")
-    return ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=vector_db.as_retriever(),
-        memory=memory,
-        return_source_documents=True
-    )
+class LLMEngine:
+    @staticmethod
+    def initialize(model_name, temp=0.7, max_tokens=512):
+        """Inicialização agressiva do modelo com otimizações de baixo nível"""
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(LLM_MODELS[model_name])
+            
+            pipe = pipeline(
+                "text-generation",
+                model=LLM_MODELS[model_name],
+                tokenizer=tokenizer,
+                device=DEVICE,
+                torch_dtype=TORCH_DTYPE,
+                max_new_tokens=max_tokens,
+                do_sample=True,
+                top_k=50,
+                temperature=temp,
+                model_kwargs={"load_in_4bit": True} if "cuda" in DEVICE else {}
+            )
+            
+            return HuggingFacePipeline(pipeline=pipe)
+        except KeyError:
+            raise ValueError("MODELO NÃO SUPORTADO!")
+        except Exception as e:
+            raise RuntimeError(f"FALHA NUCLEAR NO MODELO: {str(e)}")
 
-# Interface Gradio
-def demo():
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        vector_db = gr.State(None)
-        qa_chain = gr.State(None)
+# ===================================================================
+# INTERFACE DE COMBATE
+# ===================================================================
+def create_war_interface():
+    with gr.Blocks(theme=gr.themes.Soft(), title="⚔️ PDF Assault v1.0") as warzone:
+        state = gr.State({"db": None, "llm": None})
         
-        gr.Markdown("## 🤖 Chatbot para PDFs com Modelos Gratuitos")
+        # Zona de Upload
+        with gr.Row(variant="panel"):
+            file_upload = gr.Files(label="DOCUMENTOS ALVO", file_types=[".pdf"])
+            process_btn = gr.Button("ATAQUE!", variant="stop")
         
-        with gr.Tab("📤 Upload PDF"):
-            pdf_input = gr.Files(label="Selecione seus PDFs", file_types=[".pdf"])
-            
-        with gr.Tab("⚙️ Processamento"):
-            chunk_size = gr.Slider(100, 1000, value=500, label="Tamanho dos Chunks")
-            chunk_overlap = gr.Slider(0, 200, value=50, label="Sobreposição")
-            process_btn = gr.Button("Processar PDFs")
-            process_status = gr.Textbox(label="Status do Processamento", interactive=False)
-            
-        with gr.Tab("🧠 Modelo"):
-            model_selector = gr.Dropdown(list_llm_simple, label="Selecione o Modelo", value=list_llm_simple[1])
-            temperature = gr.Slider(0, 1, value=0.7, label="Criatividade")
-            load_model_btn = gr.Button("Carregar Modelo")
-            model_status = gr.Textbox(label="Status do Modelo", interactive=False)
-            
-        with gr.Tab("💬 Chat"):
-            chatbot = gr.Chatbot(height=400)
-            msg = gr.Textbox(label="Sua mensagem")
-            clear_btn = gr.Button("Limpar Chat")
-            
-        # Eventos
-        def process_documents(files, cs, co):
-            try:
-                file_paths = [f.name for f in files]
-                splits = load_doc(file_paths, cs, co)
-                db = create_db(splits, "docs")
-                return db, "Documentos processados!"
-            except Exception as e:
-                return None, f"Erro: {str(e)}"
+        # Controles Táticos
+        with gr.Row(variant="compact"):
+            model_selector = gr.Dropdown(list(LLM_MODELS.keys()), label="ARMA PRINCIPAL", value="TinyLlama-1.1B-Chat")
+            temp_slider = gr.Slider(0, 1, 0.7, label="POTÊNCIA DE FOGO")
+            deploy_btn = gr.Button("DEPLOY MODELO", variant="primary")
         
-        process_btn.click(
-            process_documents,
-            inputs=[pdf_input, chunk_size, chunk_overlap],
-            outputs=[vector_db, process_status]
-        )
+        # Campo de Batalha
+        chatbot = gr.Chatbot(height=600, label="ZONA DE OPERAÇÕES")
+        msg_input = gr.Textbox(label="COMANDO DE ATAQUE", placeholder="Insira o alvo...")
         
-        def load_model(model, temp, vector_db_state):
+        # Sistema de Logs
+        combat_log = gr.Textbox(label="RELATÓRIO DE COMBATE", interactive=False)
+
+        # ===== Operações Militares =====
+        @process_btn.click(inputs=[file_upload], outputs=[state, combat_log])
+        def assault_documents(files):
             try:
-                if vector_db_state is None:
-                    raise ValueError("Processe os documentos primeiro.")
-                
-                model_name = list_llm[list_llm_simple.index(model)]
-                qa = initialize_llmchain(model_name, temp, 512, 3, vector_db_state)
-                return qa, "Modelo carregado!"
+                splits = DocumentProcessor.load_and_split(files)
+                db = VectorDBManager.create(splits)
+                return {"db": db, "llm": None}, "✅ DOCUMENTOS CAPTURADOS!"
             except Exception as e:
-                return None, f"Erro: {str(e)}"
-        
-        load_model_btn.click(
-            load_model,
-            inputs=[model_selector, temperature, vector_db],
-            outputs=[qa_chain, model_status]
-        )
-        
-        def respond(message, chat_history):
-            if not qa_chain.value:
-                return "Erro: Modelo não carregado ou documentos não processados!", chat_history
+                return state.value, f"☠️ FALHA CATACLÍSMICA: {str(e)}"
+
+        @deploy_btn.click(inputs=[model_selector, temp_slider, state], outputs=[state, combat_log])
+        def deploy_model(model, temp, current_state):
+            try:
+                llm = LLMEngine.initialize(model, temp)
+                current_state["llm"] = ConversationalRetrievalChain.from_llm(
+                    llm=llm,
+                    retriever=current_state["db"].as_retriever(),
+                    memory=ConversationBufferMemory(memory_key="chat_history", return_messages=True),
+                    return_source_documents=True
+                )
+                return current_state, f"🚀 {model} PRONTO PARA COMBATE!"
+            except Exception as e:
+                return current_state, f"💥 FALHA NO DEPLOY: {str(e)}"
+
+        @msg_input.submit(inputs=[msg_input, chatbot, state], outputs=[msg_input, chatbot])
+        def execute_combat(command, history, state):
+            if not state["llm"]:
+                return command, history + [(command, "⚠️ MODELO NÃO DEPLOYADO!")]
             
             try:
-                result = qa_chain.value({"question": message, "chat_history": chat_history})
-                response = result["answer"]
-                
-                sources = "\n".join([f"📄 Página {doc.metadata['page']+1}: {doc.page_content[:50]}..." 
-                                    for doc in result.get("source_documents", [])[:2]])
-                
-                chat_history.append((message, f"{response}\n\n🔍 Fontes:\n{sources}"))
-                return "", chat_history
+                result = state["llm"]({"question": command, "chat_history": history})
+                response = f"🎯 RESPOSTA:\n{result['answer']}\n\n"
+                response += "📌 INTEL:\n" + "\n".join(
+                    f"Página {doc.metadata['page']+1}: {doc.page_content[:75]}..." 
+                    for doc in result["source_documents"][:3]
+                )
+                return "", history + [(command, response)]
             except Exception as e:
-                return f"Erro na geração: {str(e)}", chat_history
-        
-        msg.submit(respond, [msg, chatbot], [msg, chatbot])
-        clear_btn.click(lambda: [], outputs=[chatbot])
-        
-    demo.launch()
+                return command, history + [(command, f"☢️ FALHA CRÍTICA: {str(e)}")]
+
+    return warzone
 
+# ===================================================================
+# INICIALIZAÇÃO DO SISTEMA
+# ===================================================================
 if __name__ == "__main__":
-    demo()
+    interface = create_war_interface()
+    interface.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        auth=("admin", "combat123"),
+        show_error=True
+    )