DHEIVER commited on
Commit
9278ac3
·
verified ·
1 Parent(s): b9d9831

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -17
app.py CHANGED
@@ -1,17 +1,15 @@
1
  import gradio as gr
2
  import os
3
 
4
- from langchain_community.document_loaders import PyPDFLoader
5
  from langchain.text_splitter import RecursiveCharacterTextSplitter
6
- from langchain_community.vectorstores import Chroma
7
  from langchain.chains import ConversationalRetrievalChain
8
- from langchain_community.embeddings import HuggingFaceEmbeddings
9
- from langchain_community.llms import HuggingFacePipeline
10
  from langchain.memory import ConversationBufferMemory
11
 
12
  from transformers import AutoTokenizer, pipeline
13
- import torch
14
-
15
  import chromadb
16
 
17
  # Lista de modelos 100% abertos e gratuitos
@@ -32,12 +30,12 @@ list_llm_simple = [os.path.basename(llm) for llm in list_llm]
32
 
33
  # Função para carregar documentos PDF
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
  return text_splitter.split_documents(pages)
@@ -48,8 +46,7 @@ def create_db(splits, collection_name):
48
  return Chroma.from_documents(
49
  documents=splits,
50
  embedding=embedding,
51
- client=chromadb.PersistentClient(),
52
- collection_name=collection_name
53
  )
54
 
55
  # Função para inicializar o modelo LLM
@@ -63,8 +60,8 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
63
  "text-generation",
64
  model=llm_model,
65
  tokenizer=tokenizer,
66
- torch_dtype=torch.bfloat16,
67
- device_map="auto",
68
  max_new_tokens=max_tokens,
69
  do_sample=True,
70
  top_k=top_k,
@@ -90,8 +87,8 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, pr
90
  # Interface Gradio
91
  def demo():
92
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
93
- vector_db = gr.State()
94
- qa_chain = gr.State()
95
 
96
  gr.Markdown("## 🤖 Chatbot para PDFs com Modelos Gratuitos")
97
 
@@ -116,14 +113,26 @@ def demo():
116
  clear_btn = gr.Button("Limpar Chat")
117
 
118
  # Eventos
 
 
 
 
 
 
119
  process_btn.click(
120
- lambda files, cs, co: (create_db(load_doc([f.name for f in files], cs, co), "docs"), "Documentos processados!"),
121
  inputs=[pdf_input, chunk_size, chunk_overlap],
122
  outputs=[vector_db, process_status]
123
  )
124
 
 
 
 
 
 
 
125
  load_model_btn.click(
126
- lambda model, temp: (initialize_llmchain(list_llm[list_llm_simple.index(model)], temp, 512, 3, vector_db.value), "Modelo carregado!"),
127
  inputs=[model_selector, temperature],
128
  outputs=[qa_chain, model_status]
129
  )
@@ -147,4 +156,4 @@ def demo():
147
  demo.launch()
148
 
149
  if __name__ == "__main__":
150
- demo()
 
1
  import gradio as gr
2
  import os
3
 
4
+ from langchain.document_loaders import PyPDFLoader
5
  from langchain.text_splitter import RecursiveCharacterTextSplitter
6
+ from langchain.vectorstores import Chroma
7
  from langchain.chains import ConversationalRetrievalChain
8
+ from langchain.embeddings import HuggingFaceEmbeddings
9
+ from langchain.llms import HuggingFacePipeline
10
  from langchain.memory import ConversationBufferMemory
11
 
12
  from transformers import AutoTokenizer, pipeline
 
 
13
  import chromadb
14
 
15
  # Lista de modelos 100% abertos e gratuitos
 
30
 
31
  # Função para carregar documentos PDF
32
  def load_doc(list_file_path, chunk_size, chunk_overlap):
33
+ loaders = [PyPDFLoader(file_path) for file_path in list_file_path]
34
  pages = []
35
  for loader in loaders:
36
  pages.extend(loader.load())
37
  text_splitter = RecursiveCharacterTextSplitter(
38
+ chunk_size=chunk_size,
39
  chunk_overlap=chunk_overlap
40
  )
41
  return text_splitter.split_documents(pages)
 
46
  return Chroma.from_documents(
47
  documents=splits,
48
  embedding=embedding,
49
+ persist_directory=f"./{collection_name}"
 
50
  )
51
 
52
  # Função para inicializar o modelo LLM
 
60
  "text-generation",
61
  model=llm_model,
62
  tokenizer=tokenizer,
63
+ torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
64
+ device=0 if torch.cuda.is_available() else -1,
65
  max_new_tokens=max_tokens,
66
  do_sample=True,
67
  top_k=top_k,
 
87
  # Interface Gradio
88
  def demo():
89
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
90
+ vector_db = gr.State(None) # Inicializa com None
91
+ qa_chain = gr.State(None) # Inicializa com None
92
 
93
  gr.Markdown("## 🤖 Chatbot para PDFs com Modelos Gratuitos")
94
 
 
113
  clear_btn = gr.Button("Limpar Chat")
114
 
115
  # Eventos
116
+ def process_documents(files, cs, co):
117
+ file_paths = [f.name for f in files]
118
+ splits = load_doc(file_paths, cs, co)
119
+ db = create_db(splits, "docs")
120
+ return db, "Documentos processados!"
121
+
122
  process_btn.click(
123
+ process_documents,
124
  inputs=[pdf_input, chunk_size, chunk_overlap],
125
  outputs=[vector_db, process_status]
126
  )
127
 
128
+ def load_model(model, temp):
129
+ if vector_db.value is None:
130
+ return None, "Por favor, processe os documentos primeiro."
131
+ qa = initialize_llmchain(list_llm[list_llm_simple.index(model)], temp, 512, 3, vector_db.value)
132
+ return qa, "Modelo carregado!"
133
+
134
  load_model_btn.click(
135
+ load_model,
136
  inputs=[model_selector, temperature],
137
  outputs=[qa_chain, model_status]
138
  )
 
156
  demo.launch()
157
 
158
  if __name__ == "__main__":
159
+ demo()