Update app.py

app.py

@@ -1,108 +1,137 @@
-from langchain import PromptTemplate, LLMChain
-from langchain.llms import CTransformers
-import os
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
-from langchain.chains import RetrievalQA
-from langchain.embeddings import HuggingFaceBgeEmbeddings
-from io import BytesIO
-from langchain.document_loaders import PyPDFLoader
 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
 
-local_llm = "final_model_maybe_gguf-unsloth.Q5_K_M.gguf"
-
-config = {
-'max_new_tokens': 2048,
-'repetition_penalty': 1.1,
-'temperature': 0.6,
-'top_k': 50,
-'top_p': 0.9,
-'stream': True,
-'threads': int(os.cpu_count() / 2)
-}
-
-llm = CTransformers(
-    model=local_llm,
-    model_type="mistral",
-    lib="avx2", #for CPU use
-    **config
-)
-
-print("LLM Initialized...")
+token = os.environ["HF_TOKEN"]
+ST = SentenceTransformer("BM-K/KoSimCSE-roberta-multitask")
 
+dataset = load_dataset("not-lain/wikipedia",revision = "embedded")
 
-prompt_template = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
-제시된 질문에 대해서 context 내용으로 답변해줘.
+data = dataset["train"]
+data = data.add_faiss_index("embeddings") # column name that has the embeddings of the dataset
 
-### Context :
-{context}
 
-### Instruction:
-{question}
+model_id = "mintaeng/small_fut_final"
 
-### Response:
-"""
+# 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
+)
 
-model_name = "jhgan/ko-sroberta-multitask"
-model_kwargs = {'device': 'cpu'}
-encode_kwargs = {'normalize_embeddings': False}
-embeddings = HuggingFaceBgeEmbeddings(
-    model_name=model_name,
-    model_kwargs=model_kwargs,
-    encode_kwargs=encode_kwargs
+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
+
+
+@spaces.GPU(duration=150)
+def talk(prompt,history):
+    k = 1 # number of retrieved documents
+    scores , retrieved_documents = search(prompt, k)
+    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 = """
+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
+* faiss docs : https://huggingface.co/docs/datasets/v2.18.0/en/package_reference/main_classes#datasets.Dataset.add_faiss_index 
+* chatbot : https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct
+* Full documentation : https://huggingface.co/blog/not-lain/rag-chatbot-using-llama3 
+"""
 
 
-prompt = PromptTemplate(template=prompt_template, input_variables=['context', 'question'])
-load_vector_store = Chroma(persist_directory="stores/pet_cosine", embedding_function=embeddings)
-retriever = load_vector_store.as_retriever(search_kwargs={"k":1})
-# query = "what is the fastest speed for a greyhound dog?"
-# semantic_search = retriever.get_relevant_documents(query)
-# print(semantic_search)
-
-print("######################################################################")
-
-chain_type_kwargs = {"prompt": prompt}
-
-# qa = RetrievalQA.from_chain_type(
-#     llm=llm,
-#     chain_type="stuff",
-#     retriever=retriever,
-#     return_source_documents = True,
-#     chain_type_kwargs= chain_type_kwargs,
-#     verbose=True
-# )
-
-# response = qa(query)
-
-# print(response)
-
-sample_prompts = ["what is the fastest speed for a greyhound dog?", "Why should we not feed chocolates to the dogs?", "Name two factors which might contribute to why some dogs might get scared?"]
-
-def get_response(input):
-  query = input
-  chain_type_kwargs = {"prompt": prompt}
-  qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True, chain_type_kwargs=chain_type_kwargs, verbose=True)
-  response = qa(query)
-  return response
-
-input = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-iface = gr.Interface(fn=get_response, 
-             inputs=input, 
-             outputs="text",
-             title="My Dog PetCare Bot",
-             description="This is a RAG implementation based on Zephyr 7B Beta LLM.",
-             examples=sample_prompts,
-             allow_screenshot=False,
-             allow_flagging=False
-             )
-
-iface.launch()
+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's anarchy ? "]],
+    title=TITLE,
+    description=DESCRIPTION,
+    
+)
+demo.launch(debug=True)