import gradio as gr from datasets import load_dataset import os import spaces from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer import torch from threading import Thread from sentence_transformers import SentenceTransformer from datasets import load_dataset import time token = os.environ["HF_TOKEN"] model = AutoModelForCausalLM.from_pretrained( "google/gemma-7b-it", # torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, torch_dtype=torch.float16, token=token, ) tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it", token=token) device = torch.device("cuda") model = model.to(device) RAG = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") TOP_K = 1 HEADER = "\n# RESOURCES:\n" # prepare data # since data is too big we will only select the first 3K lines data = load_dataset("not-lain/wikipedia-small-3000-embedded", split="train") # index dataset data.add_faiss_index("embedding") def search(query: str, k: int = TOP_K): embedded_query = RAG.encode(query) scores, retrieved_examples = data.get_nearest_examples( "embedding", embedded_query, k=k ) return retrieved_examples def prepare_prompt(query, retrieved_examples): prompt = ( f"Query: {query}\nContinue to answer the query in short sentences by using the Search Results:\n" ) urls = [] titles = retrieved_examples["title"][::-1] texts = retrieved_examples["text"][::-1] urls = retrieved_examples["url"][::-1] titles = titles[::-1] for i in range(TOP_K): prompt += f"* {texts[i]}\n" return prompt, zip(titles, urls) @spaces.GPU(duration=150) def talk(message, history): print("history, ", history) print("message ", message) print("searching dataset ...") retrieved_examples = search(message) print("preparing prompt ...") message, metadata = prepare_prompt(message, retrieved_examples) resources = HEADER print("preparing metadata ...") for title, url in metadata: resources += f"[{title}]({url}), " print("preparing chat template ...") chat = [] for item in history: chat.append({"role": "user", "content": item[0]}) cleaned_past = item[1].split(HEADER)[0] chat.append({"role": "assistant", "content": cleaned_past}) chat.append({"role": "user", "content": message}) messages = tokenizer.apply_chat_template( chat, tokenize=False, add_generation_prompt=True ) print("chat template prepared, ", messages) print("tokenizing input ...") # Tokenize the messages string model_inputs = tokenizer([messages], return_tensors="pt").to(device) streamer = TextIteratorStreamer( tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True ) generate_kwargs = dict( model_inputs, streamer=streamer, max_new_tokens=1024, do_sample=True, top_p=0.95, top_k=1000, temperature=0.75, num_beams=1, ) print("initializing thread ...") t = Thread(target=model.generate, kwargs=generate_kwargs) t.start() time.sleep(1) # Initialize an empty string to store the generated text partial_text = "" i = 0 while t.is_alive(): try: for new_text in streamer: if new_text is not None: partial_text += new_text yield partial_text except Exception as e: print(f"retry number {i}\n LOGS:\n") i+=1 print(e, e.args) partial_text += resources yield partial_text TITLE = "# RAG" DESCRIPTION = """ A rag pipeline with a chatbot feature Resources used to build this project : * embedding model : https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1 * dataset : https://huggingface.co/datasets/not-lain/wikipedia-small-3000-embedded (used mxbai-colbert-large-v1 to create the embedding column ) * faiss docs : https://huggingface.co/docs/datasets/v2.18.0/en/package_reference/main_classes#datasets.Dataset.add_faiss_index * chatbot : https://huggingface.co/google/gemma-7b-it If you want to support my work consider clicking on the heart react button ❤️🤗 """ 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, ), examples=[["what's anarchy ? "]], title=TITLE, description=DESCRIPTION, ) demo.launch(debug=True)