Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
from datasets import load_dataset | |
import os | |
import spaces | |
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig | |
import torch | |
from threading import Thread | |
from sentence_transformers import SentenceTransformer | |
from datasets import load_dataset | |
import time | |
token = os.environ["HF_TOKEN"] | |
ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") | |
dataset = load_dataset("AI-4-Health/embedded-dataset") | |
data = dataset["train"] | |
data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset | |
model_id = "meta-llama/Meta-Llama-3-8B-Instruct" | |
# use quantization to lower GPU usage | |
bnb_config = BitsAndBytesConfig( | |
load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16 | |
) | |
tokenizer = AutoTokenizer.from_pretrained(model_id,token=token) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_id, | |
torch_dtype=torch.bfloat16, | |
device_map="auto", | |
quantization_config=bnb_config, | |
token=token | |
) | |
terminators = [ | |
tokenizer.eos_token_id, | |
tokenizer.convert_tokens_to_ids("<|eot_id|>") | |
] | |
SYS_PROMPT = """You are an assistant for answering questions. | |
You are given the extracted parts of a long document and a question. Provide a conversational answer. | |
If you don't know the answer, just say "I do not know." Don't make up an answer.""" | |
def search(query: str, k: int = 3 ): | |
"""a function that embeds a new query and returns the most probable results""" | |
embedded_query = ST.encode(query) # embed new query | |
scores, retrieved_examples = data.get_nearest_examples( # retrieve results | |
"embeddings", embedded_query, # compare our new embedded query with the dataset embeddings | |
k=k # get only top k results | |
) | |
return scores, retrieved_examples | |
def format_prompt(prompt,retrieved_documents,k): | |
"""using the retrieved documents we will prompt the model to generate our responses""" | |
PROMPT = f"Question:{prompt}\nContext:" | |
for idx in range(k) : | |
PROMPT+= f"{retrieved_documents['text'][idx]}\n" | |
return PROMPT | |
def talk(prompt,history): | |
k = 1 # number of retrieved documents | |
scores , retrieved_documents = search(prompt, k) | |
print(retrieved_documents) | |
formatted_prompt = format_prompt(prompt,retrieved_documents,k) | |
formatted_prompt = formatted_prompt[:2000] # to avoid GPU OOM | |
messages = [{"role":"system","content":SYS_PROMPT},{"role":"user","content":formatted_prompt}] | |
# tell the model to generate | |
input_ids = tokenizer.apply_chat_template( | |
messages, | |
add_generation_prompt=True, | |
return_tensors="pt" | |
).to(model.device) | |
outputs = model.generate( | |
input_ids, | |
max_new_tokens=1024, | |
eos_token_id=terminators, | |
do_sample=True, | |
temperature=0.6, | |
top_p=0.9, | |
) | |
streamer = TextIteratorStreamer( | |
tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True | |
) | |
generate_kwargs = dict( | |
input_ids= input_ids, | |
streamer=streamer, | |
max_new_tokens=1024, | |
do_sample=True, | |
top_p=0.95, | |
temperature=0.75, | |
eos_token_id=terminators, | |
) | |
t = Thread(target=model.generate, kwargs=generate_kwargs) | |
t.start() | |
outputs = [] | |
for text in streamer: | |
outputs.append(text) | |
print(outputs) | |
yield "".join(outputs) | |
TITLE = "# RAG" | |
DESCRIPTION = """ | |
HPP Chatbot | |
""" | |
demo = gr.ChatInterface( | |
fn=talk, | |
chatbot=gr.Chatbot( | |
show_label=True, | |
show_share_button=True, | |
show_copy_button=True, | |
likeable=True, | |
layout="bubble", | |
bubble_full_width=False, | |
), | |
theme="Soft", | |
examples=[["what is HPP ?"]], | |
title=TITLE, | |
description=DESCRIPTION, | |
) | |
demo.launch(debug=True) | |