DHEIVER commited on
Commit
0b1845f
·
verified ·
1 Parent(s): 6b7ae1b

Delete rag_functions.py

Browse files
Files changed (1) hide show
  1. rag_functions.py +0 -193
rag_functions.py DELETED
@@ -1,193 +0,0 @@
1
- # rag_functions.py
2
- import os
3
- from langchain_community.document_loaders import PyPDFLoader
4
- from langchain.text_splitter import RecursiveCharacterTextSplitter
5
- from langchain_community.vectorstores import Chroma
6
- from langchain.chains import ConversationalRetrievalChain
7
- from langchain_community.embeddings import HuggingFaceEmbeddings
8
- from langchain_community.llms import HuggingFaceEndpoint
9
- from langchain.memory import ConversationBufferMemory
10
- from pathlib import Path
11
- import chromadb
12
- from unidecode import unidecode
13
- import re
14
-
15
- # Lista de modelos LLM disponíveis
16
- list_llm = [
17
- "mistralai/Mistral-7B-Instruct-v0.2",
18
- "mistralai/Mixtral-8x7B-Instruct-v0.1",
19
- "mistralai/Mistral-7B-Instruct-v0.1",
20
- "google/gemma-7b-it",
21
- "google/gemma-2b-it",
22
- "HuggingFaceH4/zephyr-7b-beta",
23
- "HuggingFaceH4/zephyr-7b-gemma-v0.1",
24
- "meta-llama/Llama-2-7b-chat-hf",
25
- "microsoft/phi-2",
26
- "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
27
- "mosaicml/mpt-7b-instruct",
28
- "tiiuae/falcon-7b-instruct",
29
- "google/flan-t5-xxl"
30
- ]
31
- list_llm_simple = [os.path.basename(llm) for llm in list_llm]
32
-
33
- # Função para carregar documentos PDF e dividir em chunks
34
- def load_doc(list_file_path, chunk_size, chunk_overlap):
35
- loaders = [PyPDFLoader(x) for x in list_file_path]
36
- pages = []
37
- for loader in loaders:
38
- pages.extend(loader.load())
39
- text_splitter = RecursiveCharacterTextSplitter(
40
- chunk_size=chunk_size,
41
- chunk_overlap=chunk_overlap
42
- )
43
- doc_splits = text_splitter.split_documents(pages)
44
- return doc_splits
45
-
46
- # Função para criar o banco de dados vetorial
47
- def create_db(splits, collection_name):
48
- embedding = HuggingFaceEmbeddings()
49
- # Usando PersistentClient para persistir o banco de dados
50
- new_client = chromadb.PersistentClient(path="./chroma_db")
51
- vectordb = Chroma.from_documents(
52
- documents=splits,
53
- embedding=embedding,
54
- client=new_client,
55
- collection_name=collection_name,
56
- )
57
- return vectordb
58
-
59
- # Função para inicializar a cadeia de QA com o modelo LLM
60
- def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=None):
61
- if progress:
62
- progress(0.1, desc="Inicializando tokenizer da HF...")
63
- progress(0.5, desc="Inicializando Hub da HF...")
64
- if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
65
- llm = HuggingFaceEndpoint(
66
- repo_id=llm_model,
67
- temperature=temperature,
68
- max_new_tokens=max_tokens,
69
- top_k=top_k,
70
- load_in_8bit=True,
71
- )
72
- elif llm_model in ["HuggingFaceH4/zephyr-7b-gemma-v0.1", "mosaicml/mpt-7b-instruct"]:
73
- raise ValueError("O modelo LLM é muito grande para ser carregado automaticamente no endpoint de inferência gratuito")
74
- elif llm_model == "microsoft/phi-2":
75
- llm = HuggingFaceEndpoint(
76
- repo_id=llm_model,
77
- temperature=temperature,
78
- max_new_tokens=max_tokens,
79
- top_k=top_k,
80
- trust_remote_code=True,
81
- torch_dtype="auto",
82
- )
83
- elif llm_model == "TinyLlama/TinyLlama-1.1B-Chat-v1.0":
84
- llm = HuggingFaceEndpoint(
85
- repo_id=llm_model,
86
- temperature=temperature,
87
- max_new_tokens=250,
88
- top_k=top_k,
89
- )
90
- elif llm_model == "meta-llama/Llama-2-7b-chat-hf":
91
- raise ValueError("O modelo Llama-2-7b-chat-hf requer uma assinatura Pro...")
92
- else:
93
- llm = HuggingFaceEndpoint(
94
- repo_id=llm_model,
95
- temperature=temperature,
96
- max_new_tokens=max_tokens,
97
- top_k=top_k,
98
- )
99
-
100
- if progress:
101
- progress(0.75, desc="Definindo memória de buffer...")
102
- memory = ConversationBufferMemory(
103
- memory_key="chat_history",
104
- output_key='answer',
105
- return_messages=True
106
- )
107
- retriever = vector_db.as_retriever()
108
- if progress:
109
- progress(0.8, desc="Definindo cadeia de recuperação...")
110
- qa_chain = ConversationalRetrievalChain.from_llm(
111
- llm,
112
- retriever=retriever,
113
- chain_type="stuff",
114
- memory=memory,
115
- return_source_documents=True,
116
- verbose=False,
117
- )
118
- if progress:
119
- progress(0.9, desc="Concluído!")
120
- return qa_chain
121
-
122
- # Função para gerar um nome de coleção válido
123
- def create_collection_name(filepath):
124
- collection_name = Path(filepath).stem
125
- collection_name = collection_name.replace(" ", "-")
126
- collection_name = unidecode(collection_name)
127
- collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name)
128
- collection_name = collection_name[:50]
129
- if len(collection_name) < 3:
130
- collection_name = collection_name + 'xyz'
131
- if not collection_name[0].isalnum():
132
- collection_name = 'A' + collection_name[1:]
133
- if not collection_name[-1].isalnum():
134
- collection_name = collection_name[:-1] + 'Z'
135
- print('Caminho do arquivo: ', filepath)
136
- print('Nome da coleção: ', collection_name)
137
- return collection_name
138
-
139
- # Função para inicializar o banco de dados
140
- def initialize_database(list_file_obj, chunk_size, chunk_overlap, progress=None):
141
- list_file_path = [x.name for x in list_file_obj if x is not None]
142
- if progress:
143
- progress(0.1, desc="Criando nome da coleção...")
144
- collection_name = create_collection_name(list_file_path[0])
145
- if progress:
146
- progress(0.25, desc="Carregando documento...")
147
- doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
148
- if progress:
149
- progress(0.5, desc="Gerando banco de dados vetorial...")
150
- vector_db = create_db(doc_splits, collection_name)
151
- if progress:
152
- progress(0.9, desc="Concluído!")
153
- return vector_db, collection_name, "Completo!"
154
-
155
- # Função para inicializar o modelo LLM
156
- def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=None):
157
- llm_name = list_llm[llm_option]
158
- print("Nome do LLM: ", llm_name)
159
- qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
160
- return qa_chain, "Completo!"
161
-
162
- # Função para formatar o histórico de conversa
163
- def format_chat_history(message, chat_history):
164
- formatted_chat_history = []
165
- for user_message, bot_message in chat_history:
166
- formatted_chat_history.append(f"Usuário: {user_message}")
167
- formatted_chat_history.append(f"Assistente: {bot_message}")
168
- return formatted_chat_history
169
-
170
- # Função para realizar a conversa com o chatbot
171
- def conversation(qa_chain, message, history):
172
- formatted_chat_history = format_chat_history(message, history)
173
- response = qa_chain({"question": message, "chat_history": formatted_chat_history})
174
- response_answer = response["answer"]
175
- if response_answer.find("Resposta útil:") != -1:
176
- response_answer = response_answer.split("Resposta útil:")[-1]
177
- response_sources = response["source_documents"]
178
- response_source1 = response_sources[0].page_content.strip()
179
- response_source2 = response_sources[1].page_content.strip()
180
- response_source3 = response_sources[2].page_content.strip()
181
- response_source1_page = response_sources[0].metadata["page"] + 1
182
- response_source2_page = response_sources[1].metadata["page"] + 1
183
- response_source3_page = response_sources[2].metadata["page"] + 1
184
- new_history = history + [(message, response_answer)]
185
- return qa_chain, "", new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
186
-
187
- # Função para carregar arquivos
188
- def upload_file(file_obj):
189
- list_file_path = []
190
- for idx, file in enumerate(file_obj):
191
- file_path = file_obj.name
192
- list_file_path.append(file_path)
193
- return list_file_path