Spaces:
No application file
No application file
File size: 3,183 Bytes
0aefb9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """
"""
from langchain.document_loaders import UnstructuredFileLoader, TextLoader
from langchain.document_loaders import UnstructuredURLLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter, CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
# INSTALL CHROMADB pip install chromadb
from langchain.vectorstores import Chroma
from langchain.chains.question_answering import load_qa_chain
from langchain import OpenAI, VectorDBQA
from langchain.prompts import PromptTemplate
import os
# Defina a sua chave OpenAI:
#import openai
import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file
#openai.api_key = os.environ['OPENAI_API_KEY']
# Definimos um diretório persist para o chromaDB para que DB seja armazenada LOCALMENTE:
persist_directory = 'db'
embeddings = OpenAIEmbeddings()
# Verifique se o banco de dados já existe:
if not os.path.exists(persist_directory):
# load book.txt encode and decode in ascii and write to file to avoid encoding errors
with open("book.txt", "r",encoding="utf-8") as f:
book = f.read().encode("ascii", "ignore").decode("ascii")
with open("book_ascii.txt", "w") as f:
f.write(book)
print("Loading book_ascii.txt")
# loader = UnstructuredFileLoader("book_ascii.txt")
loader = TextLoader('book_ascii.txt')
documents = loader.load()
# print(docs[0].page_content[:])
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)
# print(texts[0].page_content)
print("Embedding book_ascii.txt")
db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
else:
db = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
# CUSTOM PROMPT
prompt_template = """Use os seguintes trechos de contexto para responder à pergunta no final,
resumindo o contexto. Se você não sabe a resposta, apenas diga que não sabe,
não tente inventar uma resposta.
{context}
Pergunta: {question}
Resposta:
"""
PROMPT = PromptTemplate(template=prompt_template,
input_variables=["context", "question"]
)
chain_type_kwargs = {"prompt": PROMPT}
# Aqui estamos usando um Banco de Vetores como um BANCO DE DADOS e não documentos pesquisados por similaridade, pois isso é feito na Chain:
qa = VectorDBQA.from_chain_type(llm=OpenAI(temperature=0.0),
chain_type="stuff",
vectorstore=db,
return_source_documents=False,
chain_type_kwargs=chain_type_kwargs
)
from langchain.callbacks import get_openai_callback
while True:
with get_openai_callback() as cb:
query = input("Digite sua query: ")
result = qa(query)
print(result['result'])
# print(result['source_documents'])
print("Tokens usados: ", cb.total_tokens) |