Update app.py

app.py

@@ -1,117 +1,8 @@
 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
+    return "hi"
 
 
 TITLE = "# RAG"
@@ -127,6 +18,8 @@ Resources used to build this project :
 * chatbot : https://huggingface.co/google/gemma-7b-it
 
 If you want to support my work consider clicking on the heart react button ❤️🤗
+
+(testing the ui)
 """