Update app.py

app.py

@@ -26,7 +26,7 @@ list_llm = [
     "mosaicml/mpt-7b-instruct"
 ]
 
-list_llm_simple = [os.path.basename(llm) for llm in list_llm]
+list_llm_simple = [name.split("/")[-1] for name in list_llm]
 
 # Função para carregar documentos PDF
 def load_doc(list_file_path, chunk_size, chunk_overlap):
@@ -56,12 +56,21 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
     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(
-        "text-generation",
+        task,
         model=llm_model,
         tokenizer=tokenizer,
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        device=0 if torch.cuda.is_available() else -1,
+        device=device,
         max_new_tokens=max_tokens,
         do_sample=True,
         top_k=top_k,
@@ -87,8 +96,8 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
 # Interface Gradio
 def demo():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        vector_db = gr.State(None)  # Inicializa com None
-        qa_chain = gr.State(None)  # Inicializa com None
+        vector_db = gr.State(None)
+        qa_chain = gr.State(None)
         
         gr.Markdown("## 🤖 Chatbot para PDFs com Modelos Gratuitos")
         
@@ -102,7 +111,7 @@ def demo():
             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[0])
+            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)
@@ -114,10 +123,13 @@ def demo():
             
         # Eventos
         def process_documents(files, cs, co):
-            file_paths = [f.name for f in files]
-            splits = load_doc(file_paths, cs, co)
-            db = create_db(splits, "docs")
-            return db, "Documentos processados!"
+            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)}"
         
         process_btn.click(
             process_documents,
@@ -126,10 +138,15 @@ def demo():
         )
         
         def load_model(model, temp, vector_db_state):
-            if vector_db_state is None:
-                return None, "Por favor, processe os documentos primeiro."
-            qa = initialize_llmchain(list_llm[list_llm_simple.index(model)], temp, 512, 3, vector_db_state)
-            return qa, "Modelo carregado!"
+            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!"
+            except Exception as e:
+                return None, f"Erro: {str(e)}"
         
         load_model_btn.click(
             load_model,
@@ -138,17 +155,20 @@ def demo():
         )
         
         def respond(message, chat_history):
-            if qa_chain.value is None:
-                return "Por favor, carregue um modelo primeiro.", chat_history
-            
-            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]])
+            if not qa_chain.value:
+                return "Erro: Modelo não carregado ou documentos não processados!", chat_history
             
-            chat_history.append((message, f"{response}\n\n🔍 Fontes:\n{sources}"))
-            return "", chat_history
+            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
+            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])