Spaces:
Sleeping
Sleeping
import time | |
import torch | |
from auto_gptq import AutoGPTQForCausalLM | |
from huggingface_hub import hf_hub_download | |
from langchain.chains import RetrievalQA | |
from langchain.embeddings import HuggingFaceInstructEmbeddings | |
from langchain.llms import HuggingFacePipeline, LlamaCpp | |
from langchain.prompts import PromptTemplate | |
from langchain.vectorstores import Chroma | |
from transformers import ( | |
AutoModelForCausalLM, | |
AutoTokenizer, | |
GenerationConfig, | |
pipeline, | |
) | |
from constants import CHROMA_SETTINGS, EMBEDDING_MODEL_NAME, PERSIST_DIRECTORY | |
def load_model(model_id, model_basename=None): | |
if model_basename is not None: | |
if ".ggml" in model_basename: | |
model_path = hf_hub_download( | |
repo_id=model_id, | |
filename=model_basename | |
) | |
max_ctx_size = 2048 | |
kwargs = {"model_path": model_path, | |
"n_ctx": max_ctx_size, | |
"max_tokens": max_ctx_size, | |
"n_gpu_layers": 1000, | |
"n_batch": max_ctx_size} | |
return LlamaCpp(**kwargs) | |
else: | |
if ".safetensors" in model_basename: | |
model_basename = model_basename.replace(".safetensors", "") | |
tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True) | |
model = AutoGPTQForCausalLM.from_quantized( | |
model_id, | |
model_basename=model_basename, | |
use_safetensors=True, | |
trust_remote_code=True, | |
device="cuda:0", | |
use_triton=False, | |
quantize_config=None, | |
) | |
else: | |
tokenizer = AutoTokenizer.from_pretrained(model_id) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_id, | |
device_map="auto", | |
torch_dtype=torch.float16, | |
low_cpu_mem_usage=True, | |
trust_remote_code=True, | |
) | |
model.tie_weights() | |
generation_config = GenerationConfig.from_pretrained(model_id) | |
pipe = pipeline( | |
"text-generation", | |
model=model, | |
tokenizer=tokenizer, | |
min_length=512, | |
max_length=2048, | |
temperature=0, | |
top_p=0.95, | |
repetition_penalty=1.15, | |
generation_config=generation_config, | |
) | |
return HuggingFacePipeline(pipeline=pipe) | |
def load_qa(): | |
embeddings = HuggingFaceInstructEmbeddings( | |
model_name=EMBEDDING_MODEL_NAME, | |
model_kwargs={"device": "cuda"} | |
) | |
db = Chroma( | |
persist_directory=PERSIST_DIRECTORY, | |
embedding_function=embeddings, | |
client_settings=CHROMA_SETTINGS, | |
) | |
retriever = db.as_retriever() | |
model_id = "psmathur/orca_mini_3b" | |
model_basename = None | |
# model_id = "TheBloke/vicuna-7B-1.1-HF" | |
# model_basename = None | |
template = """You are an AI assistant for answering questions about logotherapy. You are given the following | |
extracted parts of a annual academic journal. Provide a very detailed comprehensive academic answer. If you don't | |
know the answer, just say "I'm not sure." Don't try to make up an answer. If the question is not about the | |
psychotherapy and not directly in the given context, politely inform them that you are tuned to only answer | |
questions about logotherapy. Question: {question} ========= {context} ========= Answer:""" | |
prompt = PromptTemplate( | |
input_variables=["context", "question"], | |
template=template | |
) | |
llm = load_model( | |
model_id=model_id, | |
model_basename=model_basename | |
) | |
qa = RetrievalQA.from_chain_type( | |
llm=llm, | |
chain_type="stuff", | |
retriever=retriever, | |
return_source_documents=True, | |
chain_type_kwargs={"prompt": prompt}, | |
) | |
return qa |