profdanielvieira95 commited on
Commit
5d91cef
·
verified ·
1 Parent(s): 3d1e57d

update app

Browse files
Files changed (1) hide show
  1. app.py +132 -123
app.py CHANGED
@@ -1,123 +1,132 @@
1
- import os
2
- import gradio as gr
3
- from typing import List
4
- from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
5
- from llama_index.core.node_parser import SentenceSplitter
6
- from llama_index.embeddings.huggingface import HuggingFaceEmbedding
7
- from llama_index.vector_stores.chroma import ChromaVectorStore
8
- from llama_index.llms.groq import Groq
9
- from llama_index.core.memory import ChatSummaryMemoryBuffer
10
- import chromadb
11
- from tempfile import TemporaryDirectory
12
- from PyPDF2 import PdfReader
13
- from corretor import corrigir_texto # <<< Correção importada aqui
14
-
15
- # Wrapper de embedding compatível com ChromaDB
16
- class ChromaEmbeddingWrapper:
17
- def __init__(self, model_name: str):
18
- self.model = HuggingFaceEmbedding(model_name=model_name)
19
-
20
- def __call__(self, input: List[str]) -> List[List[float]]:
21
- return self.model.embed_documents(input)
22
-
23
- # Inicializa modelos de embedding
24
- embed_model = HuggingFaceEmbedding(model_name='intfloat/multilingual-e5-large')
25
- embed_model_chroma = ChromaEmbeddingWrapper(model_name='intfloat/multilingual-e5-large')
26
-
27
- # Inicializa ChromaDB
28
- chroma_client = chromadb.PersistentClient(path='./chroma_db')
29
- collection_name = 'documentos_bitdoglab'
30
- chroma_collection = chroma_client.get_or_create_collection(
31
- name=collection_name,
32
- embedding_function=embed_model_chroma
33
- )
34
-
35
- vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
36
- storage_context = StorageContext.from_defaults(vector_store=vector_store)
37
-
38
- # Inicializa LLM da Groq
39
- Groq_api = os.environ.get("GROQ_API_KEY")
40
- llms = Groq(model='llama3-70b-8192', api_key=Groq_api or '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')
41
-
42
- # Estados globais
43
- document_index = None
44
- chat_engine = None
45
-
46
- # Carregamento único do PDF
47
- def carregar_pdf_inicial():
48
- global document_index, chat_engine
49
-
50
- try:
51
- with TemporaryDirectory() as tmpdir:
52
- pdf_path = "BitDogLab_info_v2.pdf"
53
- text = ""
54
- reader = PdfReader(pdf_path)
55
- for page in reader.pages:
56
- text += page.extract_text() or ""
57
-
58
- with open(os.path.join(tmpdir, "temp.txt"), "w", encoding="utf-8") as f:
59
- f.write(text)
60
-
61
- documentos = SimpleDirectoryReader(input_dir=tmpdir)
62
- docs = documentos.load_data()
63
-
64
- node_parser = SentenceSplitter(chunk_size=1200,chunk_overlap=150)
65
- nodes = node_parser.get_nodes_from_documents(docs, show_progress=True)
66
-
67
- document_index = VectorStoreIndex(nodes, storage_context=storage_context, embed_model=embed_model)
68
-
69
- memory = ChatSummaryMemoryBuffer(llm=llms, token_limit=256)
70
-
71
- chat_engine = document_index.as_chat_engine(
72
- chat_mode='context',
73
- llm=llms,
74
- memory=memory,
75
- system_prompt='''Você é especialista na placa BitDog Lab e sua função é ajudar os usuários nas dúvidas e informações sobre a placa e como criar códigos.'''
76
- )
77
-
78
- print("PDF carregado com sucesso.")
79
-
80
- except Exception as e:
81
- print(f"Erro ao carregar PDF: {e}")
82
-
83
- # Função de chat com correção de texto
84
- def converse_com_bot(message, chat_history):
85
- global chat_engine
86
-
87
- if chat_engine is None:
88
- return "Erro: o bot ainda não está pronto.", chat_history
89
-
90
- response = chat_engine.chat(message)
91
- resposta_corrigida = corrigir_texto(response.response) # <<< Aplica correção
92
-
93
- if chat_history is None:
94
- chat_history = []
95
-
96
- chat_history.append({"role": "user", "content": message})
97
- chat_history.append({"role": "assistant", "content": resposta_corrigida})
98
-
99
- return "", chat_history
100
-
101
- # Resetar conversa
102
- def resetar_chat():
103
- global chat_engine
104
- if chat_engine:
105
- chat_engine.reset()
106
- return []
107
-
108
- # Carregar PDF na inicialização
109
- carregar_pdf_inicial()
110
-
111
- # Interface Gradio
112
- with gr.Blocks() as app:
113
- gr.Markdown("# 🤖 Chatbot BitDog Lab - Seu assistente para esclarecer dúvidas")
114
-
115
- chatbot = gr.Chatbot(label="Conversa", type="messages")
116
- msg = gr.Textbox(label='Digite a sua mensagem')
117
- limpar = gr.Button('Limpar')
118
-
119
- msg.submit(converse_com_bot, [msg, chatbot], [msg, chatbot])
120
- limpar.click(resetar_chat, None, chatbot, queue=False)
121
-
122
- #app.launch()
123
- app.launch(server_name="0.0.0.0", server_port=7860)
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ from typing import List
4
+ from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
5
+ from llama_index.core.node_parser import SentenceSplitter
6
+ from llama_index.embeddings.huggingface import HuggingFaceEmbedding
7
+ from llama_index.vector_stores.chroma import ChromaVectorStore
8
+ from llama_index.llms.groq import Groq
9
+ from llama_index.core.memory import ChatSummaryMemoryBuffer
10
+ import chromadb
11
+ from tempfile import TemporaryDirectory
12
+ from PyPDF2 import PdfReader
13
+ from corretor import corrigir_texto # <<< Correção importada aqui
14
+ import platform
15
+
16
+ # Wrapper de embedding compatível com ChromaDB
17
+ class ChromaEmbeddingWrapper:
18
+ def __init__(self, model_name: str):
19
+ self.model = HuggingFaceEmbedding(model_name=model_name)
20
+
21
+ def __call__(self, input: List[str]) -> List[List[float]]:
22
+ return self.model.embed_documents(input)
23
+
24
+ # Inicializa modelos de embedding
25
+ embed_model = HuggingFaceEmbedding(model_name='intfloat/multilingual-e5-large')
26
+ embed_model_chroma = ChromaEmbeddingWrapper(model_name='intfloat/multilingual-e5-large')
27
+
28
+ # Inicializa ChromaDB
29
+
30
+ # Define caminho seguro dependendo do sistema operacional
31
+ if platform.system() == "Windows":
32
+ chroma_path = "./chroma_db"
33
+ else:
34
+ chroma_path = "/tmp/chroma_db"
35
+
36
+ chroma_client = chromadb.PersistentClient(path=chroma_path)
37
+
38
+ collection_name = 'documentos_bitdoglab'
39
+ chroma_collection = chroma_client.get_or_create_collection(
40
+ name=collection_name,
41
+ embedding_function=embed_model_chroma
42
+ )
43
+
44
+ vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
45
+ storage_context = StorageContext.from_defaults(vector_store=vector_store)
46
+
47
+ # Inicializa LLM da Groq
48
+ Groq_api = os.environ.get("GROQ_API_KEY")
49
+ llms = Groq(model='llama3-70b-8192', api_key=Groq_api or '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')
50
+
51
+ # Estados globais
52
+ document_index = None
53
+ chat_engine = None
54
+
55
+ # Carregamento único do PDF
56
+ def carregar_pdf_inicial():
57
+ global document_index, chat_engine
58
+
59
+ try:
60
+ with TemporaryDirectory() as tmpdir:
61
+ pdf_path = "BitDogLab_info_v2.pdf"
62
+ text = ""
63
+ reader = PdfReader(pdf_path)
64
+ for page in reader.pages:
65
+ text += page.extract_text() or ""
66
+
67
+ with open(os.path.join(tmpdir, "temp.txt"), "w", encoding="utf-8") as f:
68
+ f.write(text)
69
+
70
+ documentos = SimpleDirectoryReader(input_dir=tmpdir)
71
+ docs = documentos.load_data()
72
+
73
+ node_parser = SentenceSplitter(chunk_size=1200,chunk_overlap=150)
74
+ nodes = node_parser.get_nodes_from_documents(docs, show_progress=True)
75
+
76
+ document_index = VectorStoreIndex(nodes, storage_context=storage_context, embed_model=embed_model)
77
+
78
+ memory = ChatSummaryMemoryBuffer(llm=llms, token_limit=256)
79
+
80
+ chat_engine = document_index.as_chat_engine(
81
+ chat_mode='context',
82
+ llm=llms,
83
+ memory=memory,
84
+ system_prompt='''Você é especialista na placa BitDog Lab e sua função é ajudar os usuários nas dúvidas e informações sobre a placa e como criar códigos.'''
85
+ )
86
+
87
+ print("PDF carregado com sucesso.")
88
+
89
+ except Exception as e:
90
+ print(f"Erro ao carregar PDF: {e}")
91
+
92
+ # Função de chat com correção de texto
93
+ def converse_com_bot(message, chat_history):
94
+ global chat_engine
95
+
96
+ if chat_engine is None:
97
+ return "Erro: o bot ainda não está pronto.", chat_history
98
+
99
+ response = chat_engine.chat(message)
100
+ resposta_corrigida = corrigir_texto(response.response) # <<< Aplica correção
101
+
102
+ if chat_history is None:
103
+ chat_history = []
104
+
105
+ chat_history.append({"role": "user", "content": message})
106
+ chat_history.append({"role": "assistant", "content": resposta_corrigida})
107
+
108
+ return "", chat_history
109
+
110
+ # Resetar conversa
111
+ def resetar_chat():
112
+ global chat_engine
113
+ if chat_engine:
114
+ chat_engine.reset()
115
+ return []
116
+
117
+ # Carregar PDF na inicialização
118
+ carregar_pdf_inicial()
119
+
120
+ # Interface Gradio
121
+ with gr.Blocks() as app:
122
+ gr.Markdown("# 🤖 Chatbot BitDog Lab - Seu assistente para esclarecer dúvidas")
123
+
124
+ chatbot = gr.Chatbot(label="Conversa", type="messages")
125
+ msg = gr.Textbox(label='Digite a sua mensagem')
126
+ limpar = gr.Button('Limpar')
127
+
128
+ msg.submit(converse_com_bot, [msg, chatbot], [msg, chatbot])
129
+ limpar.click(resetar_chat, None, chatbot, queue=False)
130
+
131
+ #app.launch()
132
+ app.launch(server_name="0.0.0.0", server_port=7860)