Spaces:
Sleeping
Sleeping
from langchain_community.llms import CTransformers | |
from langchain.prompts import PromptTemplate | |
from langchain_core.runnables import RunnableSequence | |
from langchain.chains import RetrievalQA | |
#from langchain_community.embeddings import GPT4AllEmbeddings | |
from langchain_community.embeddings import HuggingFaceEmbeddings | |
from langchain_community.vectorstores import FAISS | |
from huggingface_hub import hf_hub_download | |
# !pip install llama-cpp-python | |
# from llama_cpp import Llama | |
# model_file = Llama.from_pretrained( | |
# repo_id="Pudding48/TinyLLamaTest", | |
# filename="tinyllama-1.1b-chat-v1.0.Q8_0.gguf", | |
# ) | |
import os | |
cache_path = "/home/user/app/hf_cache" | |
os.makedirs(cache_path, exist_ok=True) | |
model_file = hf_hub_download( | |
repo_id="Pudding48/TinyLlamaTest", | |
filename="tinyllama-1.1b-chat-v1.0.Q8_0.gguf", | |
cache_dir=cache_path | |
) | |
# Vector store location | |
vector_dp_path = "/home/user/app/vectorstores/db_faiss" | |
from prepare_vector_dp import create_db_from_text | |
create_db_from_text() | |
# Load LLM with CTransformers | |
def load_llm(model_file): | |
return CTransformers( | |
model=model_file, | |
model_type="llama", | |
temperature=0.01, | |
config={'gpu_layers': 0}, | |
max_new_tokens=128, | |
context_length=512 | |
) | |
# Create the prompt | |
def creat_prompt(template): | |
return PromptTemplate(template=template, input_variables=["context", "question"]) | |
# Create QA pipeline | |
def create_qa_chain(prompt, llm, db): | |
return RetrievalQA.from_chain_type( | |
llm=llm, | |
chain_type="stuff", | |
retriever=db.as_retriever(search_kwargs={"k": 1}), | |
return_source_documents=False, | |
chain_type_kwargs={'prompt': prompt} | |
) | |
# Load vector DB | |
def read_vector_db(): | |
embedding_model = HuggingFaceEmbeddings(model_name = "sentence-transformers/all-MiniLM-L6-v2") | |
return FAISS.load_local(vector_dp_path, embedding_model, allow_dangerous_deserialization=True) | |
# Build everything | |
db = read_vector_db() | |
llm = load_llm(model_file) | |
template = """<|im_start|>system\nSử dụng thông tin sau đây để trả lời câu hỏi. Nếu bạn không biết câu trả lời, hãy nói không biết, đừng cố tạo ra câu trả lời\n | |
{context}<|im_end|>\n<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant""" | |
prompt = creat_prompt(template) | |
llm_chain = create_qa_chain(prompt, llm, db) |