Spaces:
Runtime error
Runtime error
File size: 3,694 Bytes
f32218b b140cfb f32218b b140cfb f32218b b140cfb f32218b b140cfb |
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 89 90 91 92 93 94 95 96 97 |
"""Adopted from https://github.com/imartinez/privateGPT/blob/main/privateGPT.py
https://raw.githubusercontent.com/imartinez/privateGPT/main/requirements.txt
from pathlib import Path
Path("models").mkdir(exit_ok=True)
!time wget -c https://gpt4all.io/models/ggml-gpt4all-j-v1.3-groovy.bin -O models/ggml-gpt4all-j-v1.3-groovy.bin"""
from dotenv import load_dotenv, dotenv_values
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.vectorstores import Chroma
from langchain.llms import GPT4All, LlamaCpp
import os
import argparse
import time
from types import SimpleNamespace
from chromadb.config import Settings
# embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME")
# persist_directory = os.environ.get('PERSIST_DIRECTORY')
# load_dotenv()
# model_type = os.environ.get('MODEL_TYPE')
# model_path = os.environ.get('MODEL_PATH')
# model_n_ctx = os.environ.get('MODEL_N_CTX')
# model_n_batch = int(os.environ.get('MODEL_N_BATCH',8))
# target_source_chunks = int(os.environ.get('TARGET_SOURCE_CHUNKS',4))
settings = dict([('PERSIST_DIRECTORY', 'db1'),
('MODEL_TYPE', 'GPT4All'),
('MODEL_PATH', 'models/ggml-gpt4all-j-v1.3-groovy.bin'),
('EMBEDDINGS_MODEL_NAME', 'all-MiniLM-L6-v2'),
('MODEL_N_CTX', '1000'),
('MODEL_N_BATCH', '8'),
('TARGET_SOURCE_CHUNKS', '4')])
# models/ggml-gpt4all-j-v1.3-groovy.bin ~5G
persist_directory = settings.get('PERSIST_DIRECTORY')
model_type = settings.get('MODEL_TYPE')
model_path = settings.get('MODEL_PATH')
embeddings_model_name = settings.get("EMBEDDINGS_MODEL_NAME")
# embeddings_model_name = 'all-MiniLM-L6-v2'
# embeddings_model_name = 'paraphrase-multilingual-mpnet-base-v2'
model_n_ctx = settings.get('MODEL_N_CTX')
model_n_batch = int(settings.get('MODEL_N_BATCH',8))
target_source_chunks = int(settings.get('TARGET_SOURCE_CHUNKS',4))
# Define the Chroma settings
CHROMA_SETTINGS = Settings(
chroma_db_impl='duckdb+parquet',
persist_directory=persist_directory,
anonymized_telemetry=False
)
args = SimpleNamespace(hide_source=False, mute_stream=False)
# load chroma database from db1
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
retriever = db.as_retriever(search_kwargs={"k": target_source_chunks})
# activate/deactivate the streaming StdOut callback for LLMs
callbacks = [] if args.mute_stream else [StreamingStdOutCallbackHandler()]
# Prepare the LLM
match model_type:
case "LlamaCpp":
llm = LlamaCpp(model_path=model_path, n_ctx=model_n_ctx, n_batch=model_n_batch, callbacks=callbacks, verbose=False)
case "GPT4All":
llm = GPT4All(model=model_path, n_ctx=model_n_ctx, backend='gptj', n_batch=model_n_batch, callbacks=callbacks, verbose=False)
case _default:
# raise exception if model_type is not supported
raise Exception(f"Model type {model_type} is not supported. Please choose one of the following: LlamaCpp, GPT4All")
# need about 5G RAM
qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents= not args.hide_source)
# Get the answer from the chain
query = "共产党是什么"
start = time.time()
res = qa(query)
answer, docs = res['result'], [] if args.hide_source else res['source_documents']
end = time.time()
# Print the result
print("\n\n> Question:")
print(query)
print(f"\n> Answer (took {round(end - start, 2)} s.):")
print(answer)
|