Update app.py

app.py

@@ -42,75 +42,130 @@ from langchain_core.messages import SystemMessage
 from langchain_core.prompts import HumanMessagePromptTemplate
 from langchain_core.prompts import ChatPromptTemplate
 from langchain.prompts import PromptTemplate
+import time
+
 
 print("check1")
 question = "How can I reverse Diabetes?"
 
+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."""
+
+print("check2")
 
-#template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. Use three sentences maximum. Keep the answer as concise as possible. Always say "thanks for asking!" at the end of the answer. 
-#{context}
-#Question: {question}
-#Helpful Answer:"""
 
-#QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],template=template)
 
 from transformers import AutoTokenizer
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
 
 llm_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(llm_model)
-model = AutoModelForCausalLM.from_pretrained(llm_model)
-pipe = pipeline(model = llm_model, tokenizer = tokenizer, task = "text-generation", temperature=0.5)
+tokenizer = AutoTokenizer.from_pretrained(llm_model,token=token)
+model = AutoModelForCausalLM.from_pretrained(llm_model,token=token)
+#pipe = pipeline(model = llm_model, tokenizer = tokenizer, task = "text-generation", temperature=0.5)
 
-#question = "How can I reverse diabetes?"
-#docs1 = retriever.invoke(question)
-#docs1 = retriever.similarity_search(question)
-#print(docs1[0].page_content)
+terminators = [
+    tokenizer.eos_token_id,
+    tokenizer.convert_tokens_to_ids("<|eot_id|>")
+]
 
-import pandas as pd
-#df = pd.DataFrame(docs1, columns=["text"])
-#context =  df.to_string()
-#print(context)
 
-#print(docs1)[0]['generated_text'][-1]
 
-print("check2")
+def search(query: str, k: int = 3 ):
+    """a function that embeds a new query and returns the most probable results"""
+    embedded_query = embedding_model.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
 
+print("check2A")
 
-#question = "How can I reverse diabetes?"
-print("result")
 
-print("check3")
+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
 
-messages = [
-    {
-        "role": "system",
-        "content": "You are a friendly chatbot who responds in the style of a  doctor",
-    },
-    {"role": "user", "content": "How can I reverse diabetes?"},
-]
-prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-chain = pipe(prompt, max_new_tokens=256, do_sample=True)
-print(chain[0]["generated_text"])
 
-#chain = pipe(question = question,context = "Use the following information to answer the question. Diabetes can be cured by eating apples.") 
+print("check3")
 
 
+print("check3A")
 
 
+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)
 
-#chain = pipe(question = question,context = "Use the following information to answer the question. {context}.") 
-#context = "Use the following information to answer the question. Diabetes can be cured by eating apples."
 
-print("check3A")
-#print(chain)[0]['generated_text'][-1]
 print("check3B")
 
+
+TITLE = "AI Copilot for Diabetes Patients"
+
+DESCRIPTION = ""
+
 import gradio as gr
-#ragdemo = gr.Interface.from_pipeline(chain)
 
-interface = gr.Interface.from_pipeline(chain).launch(share=True)
+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 is Diabetes? "]],
+    title=TITLE,
+    description=DESCRIPTION,
+    
+)
+demo.launch(debug=True)
 
-print("check4")
-#ragdemo.launch()
-print("check5")
+print("check4")