🌘w🌖

app.py

@@ -7,7 +7,7 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
-from ragatouille import RAGPretrainedModel
+from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
 
 
@@ -18,72 +18,48 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16,
     token=token,
 )
-tok = AutoTokenizer.from_pretrained("google/gemma-7b-it", token=token)
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it", token=token)
 device = torch.device("cuda")
 model = model.to(device)
-RAG = RAGPretrainedModel.from_pretrained("mixedbread-ai/mxbai-colbert-v1")
-
+RAG = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
+TOP_K = 3
 # prepare data
 # since data is too big we will only select the first 3K lines
 
-dataset = load_dataset(
-    "wikimedia/wikipedia", "20231101.en", split="train", streaming=True
-)
-# init data
-data = Dataset.from_dict({})
-i = 0
-for i, entry in enumerate(dataset):
-    # each entry has the following columns
-    # ['id', 'url', 'title', 'text']
-    data = data.add_item(entry)
-    if i == 3000:
-        break
-# free memory
-del dataset  # we keep data
+data = load_dataset("not-lain/wikipedia-small-3000-embedded", subset="train")
 
-# index data
-documents = data["text"]
-RAG.index(documents, index_name="wikipedia", use_faiss=True)
-# free memory
-del documents
+# index dataset
+data.add_faiss_index("embedding", device=1)
 
-def search(query, k: int = 5):
-    results = RAG.search(query, k=k)
-    # results are ordered according to their score
-    # results has the following keys
-    #
-    # {'content' : 'retrieved content'
-    # 'score' : score[float]
-    # 'rank' : "results are sorted using score and each is given a rank, also can be called place, 1 2 3 4 ..."
-    # 'document_id' : "no clue man i just got here"
-    # 'passage_id' :  "or original row number"
-    # }
-    #
-    return [result["passage_id"] for result in results]
+@spaces.GPU
+def search(query: str, k: int = TOP_K):
+    embedded_query = model.encode(query)
+    scores, retrieved_examples = data.get_nearest_examples(
+        "embedding", embedded_query, k=k
+    )
+    return retrieved_examples
 
 
-def prepare_prompt(query, indexes,data = data):
+def prepare_prompt(query, retrieved_examples):
     prompt = (
         f"Query: {query}\nContinue to answer the query by using the Search Results:\n"
     )
-    titles = []
     urls = []
-    for i in indexes:
-        title = entry["title"][i]
-        text = entry["text"][i]
-        url = entry["url"][i]
-        titles.append(title)
-        urls.append(url)
-        prompt += f"Title: {title}, Text: {text}\n"
-    return prompt, (titles,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"Title: {titles[i]}, Text: {texts[i]}\n"
+    return prompt, (titles, urls)
 
 
 @spaces.GPU
 def talk(message, history):
-    indexes = search(message)
-    message,metadata = prepare_prompt(message, indexes)
+    retrieved_examples = search(message)
+    message, metadata = prepare_prompt(message, retrieved_examples)
     resources = "\nRESOURCES:\n"
-    for title,url in metadata:
+    for title, url in metadata:
         resources += f"[{title}]({url}),  "
     chat = []
     for item in history:
@@ -92,11 +68,11 @@ def talk(message, history):
             cleaned_past = item[1].split("\nRESOURCES:\n")[0]
             chat.append({"role": "assistant", "content": cleaned_past})
     chat.append({"role": "user", "content": message})
-    messages = tok.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+    messages = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
     # Tokenize the messages string
-    model_inputs = tok([messages], return_tensors="pt").to(device)
+    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
     streamer = TextIteratorStreamer(
-        tok, timeout=10.0, skip_prompt=True, skip_special_tokens=True
+        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
     )
     generate_kwargs = dict(
         model_inputs,

requirements.txt

@@ -1,6 +1,6 @@
 spaces
 torch==2.2.0
 transformers
-ragatouille
+sentence-transformers
 faiss-gpu
 datasets