from datasets import load_dataset from datasets import Dataset from langchain.docstore.document import Document as LangchainDocument from langchain.memory import ConversationBufferMemory from sentence_transformers import SentenceTransformer import faiss import time import torch import pandas as pd from io import StringIO from transformers import AutoTokenizer from transformers import AutoModelForCausalLM from transformers import TextIteratorStreamer from threading import Thread #from huggingface_hub import InferenceClient from huggingface_hub import Repository, upload_file import os HF_TOKEN = os.getenv('HF_Token') llm_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # TheBloke/Llama-2-7B-Chat-GGML , TinyLlama/TinyLlama-1.1B-Chat-v1.0 , microsoft/Phi-3-mini-4k-instruct, health360/Healix-1.1B-V1-Chat-dDPO # TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF not working tokenizer = AutoTokenizer.from_pretrained(llm_model) # pulling tokeinzer for text generation model data = load_dataset("Namitg02/Test", split='train', streaming=False) #Returns a list of dictionaries, each representing a row in the dataset. length = len(data) #print(data[2]) embedding_model = SentenceTransformer("all-MiniLM-L6-v2") embedding_dim = embedding_model.get_sentence_embedding_dimension() # Returns dimensions of embedidng index = faiss.IndexFlatL2(embedding_dim) data.add_faiss_index("embeddings", custom_index=index) # adds an index column for the embeddings print("check1d") #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.""" # Provides context of how to answer the question #messages = [ # { # "role": "system", # "content": "You are a friendly chatbot who always responds in the style of a pirate", # }, # {"role": "user", "content": "How many helicopters can a human eat in one sitting?"}, #] print("check2") memory = ConversationBufferMemory(return_messages=True) model = AutoModelForCausalLM.from_pretrained(llm_model) # Initializing the text generation model terminators = [ tokenizer.eos_token_id, # End-of-Sequence Token that indicates where the model should consider the text sequence to be complete tokenizer.convert_tokens_to_ids("<|eot_id|>") # Converts a token strings in a single/ sequence of integer id using the vocabulary ] # indicates the end of a sequence def search(query: str, k: int = 2 ): """a function that embeds a new query and returns the most probable results""" embedded_query = embedding_model.encode(query) # create embedding of a 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 # returns scores (List[float]): the retrieval scores from either FAISS (IndexFlatL2 by default) and examples (dict) format # called by talk function that passes prompt #print(scores, retrieved_examples) print("check2A") 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['0'][idx]}\n" return PROMPT # Called by talk function to add retrieved documents to the prompt. Keeps adding text of retrieved documents to string taht are retreived print("check3") #print(PROMPT) print("check3A") def talk(prompt,history): k = 2 # number of retrieved documents scores , retrieved_documents = search(prompt, k) # get retrival scores and examples in dictionary format based on the prompt passed print(retrieved_documents.keys()) formatted_prompt = format_prompt(prompt,retrieved_documents,k) # create a new prompt using the retrieved documents print(retrieved_documents['0']) print(formatted_prompt) formatted_prompt = formatted_prompt[:600] # to avoid memory issue # print(retrieved_documents['0'][1] # print(retrieved_documents['0'][2] print(formatted_prompt) messages = [{"role":"system","content":SYS_PROMPT},{"role":"user","content":formatted_prompt}] # binding the system context and new prompt for LLM # the chat template structure should be based on text generation model format print("check3B") input_ids = tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt" ).to(model.device) # tell the model to generate # add_generation_prompt argument tells the template to add tokens that indicate the start of a bot response print("check3C") outputs = model.generate( input_ids, max_new_tokens=200, eos_token_id=terminators, do_sample=True, temperature=0.4, top_p=0.9, ) # calling the model to generate response based on message/ input # do_sample if set to True uses strategies to select the next token from the probability distribution over the entire vocabulary # temperature controls randomness. more renadomness with higher temperature # only the tokens comprising the top_p probability mass are considered for responses # This output is a data structure containing all the information returned by generate(), but that can also be used as tuple or dictionary. print("check3D") streamer = TextIteratorStreamer( tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True ) # stores print-ready text in a queue, to be used by a downstream application as an iterator. removes specail tokens in generated text. # timeout for text queue. tokenizer for decoding tokens # called by generate_kwargs print("check3E") generate_kwargs = dict( input_ids= input_ids, streamer=streamer, max_new_tokens= 100, do_sample=True, top_p=0.95, temperature=0.5, eos_token_id=terminators, ) # send additional parameters to model for generation print("check3F") t = Thread(target=model.generate, kwargs=generate_kwargs) # to process multiple instances t.start() # start a thread print("check3G") outputs = [] for text in streamer: outputs.append(text) print(outputs) yield "".join(outputs) print("check3H") pd.options.display.max_colwidth = 800 outputstring = ''.join(outputs) print(outputstring) print(prompt) print(outputstring) memory_string = { "Prompt": [prompt], "Output": [outputstring] } print(type(memory_string)) memory_panda = pd.DataFrame(memory_string) print(memory_panda.iloc[[0]]) df.loc[len(df.index)] = ['prompt', outputstring] print(memory_panda.iloc[[1]]) # datasetconversation = Dataset.from_pandas(memory_panda) TITLE = "AI Copilot for Diabetes Patients" DESCRIPTION = "I provide answers to concerns related to Diabetes" import gradio as gr # Design chatbot 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, ) # launch chatbot and calls the talk function which in turn calls other functions print("check3I") demo.launch()