# -*- coding: utf-8 -*- """RAG_using_Llama3.ipynb Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/1b-ZDo3QQ-axgm804UlHu3ohZwnoXz5L1 # install dependecies """ !pip install -q datasets sentence-transformers faiss-cpu accelerate from huggingface_hub import notebook_login notebook_login() """# embed dataset this is a slow procedure so you might consider saving your results """ from datasets import load_dataset dataset = load_dataset("KarthikaRajagopal/wikipedia-2") dataset from sentence_transformers import SentenceTransformer ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") # embed the dataset def embed(batch): # or you can combine multiple columns here, for example the title and the text information = batch["text"] return {"embeddings" : ST.encode(information)} dataset = dataset.map(embed,batched=True,batch_size=16) !pip install datasets from datasets import load_dataset dataset = load_dataset("KarthikaRajagopal/wikipedia-2",revision = "embedded") # Push it to your Hugging Face repository dataset.push_to_hub("KarthikaRajagopal/wikipedia-2", revision="embedded") from datasets import load_dataset dataset = load_dataset("KarthikaRajagopal/wikipedia-2",revision = "embedded") data = dataset["train"] data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset 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 scores , result = search("anarchy", 4 ) # search for word anarchy and get the best 4 matching values from the dataset # the lower the better scores result['title'] print(result["text"][0]) """# chatbot on top of the retrieved results""" !pip install -q datasets sentence-transformers faiss-cpu accelerate bitsandbytes from sentence_transformers import SentenceTransformer ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") from datasets import load_dataset dataset = load_dataset("KarthikaRajagopal/wikipedia-2",revision = "embedded") data = dataset["train"] data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset 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 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig import torch model_id = "meta-llama/Meta-Llama-3-8B-Instruct" 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) model = AutoModelForCausalLM.from_pretrained( model_id, torch_dtype=torch.bfloat16, device_map="auto", quantization_config=bnb_config ) 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 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 generate(formatted_prompt): 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, ) response = outputs[0][input_ids.shape[-1]:] return tokenizer.decode(response, skip_special_tokens=True) def rag_chatbot(prompt:str,k:int=2): scores , retrieved_documents = search(prompt, k) formatted_prompt = format_prompt(prompt,retrieved_documents,k) return generate(formatted_prompt) rag_chatbot("what's anarchy ?", k = 2)