| from datasets import load_dataset | |
| from langchain.docstore.document import Document as LangchainDocument | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from sentence_transformers import SentenceTransformer | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| #from langchain_community.vectorstores import faiss | |
| import faiss | |
| from langchain.prompts import PromptTemplate | |
| #from langchain.chains import ConversationalRetrievalChain | |
| #from transformers import pipeline | |
| #from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline | |
| #from langchain_core.messages import SystemMessage | |
| import time | |
| from transformers import AutoTokenizer | |
| from transformers import AutoModelForCausalLM | |
| from transformers import TextIteratorStreamer | |
| from threading import Thread | |
| #dataset = load_dataset("Namitg02/Test", split='train', streaming=False) | |
| dataset = load_dataset("not-lain/wikipedia",revision = "embedded") | |
| # Returns a list of dictionaries, each representing a row in the dataset. | |
| #print(dataset[1]) | |
| #splitter = RecursiveCharacterTextSplitter(chunk_size=150, chunk_overlap=25,separators=["\n\n"]) # ["\n\n", "\n", " ", ""]) | |
| #docs = splitter.create_documents(str(dataset)) | |
| # Returns a list of documents | |
| #print(docs) | |
| embedding_model = SentenceTransformer("all-MiniLM-L6-v2") | |
| #all-MiniLM-L6-v2, BAAI/bge-base-en-v1.5,infgrad/stella-base-en-v2, BAAI/bge-large-en-v1.5 working with default dimensions | |
| #docs_text = [doc.text for doc in docs] | |
| #embed = embedding_model.embed_documents(docs_text) | |
| #data = FAISS.from_embeddings(embed, embedding_model) | |
| #data = FAISS.from_texts(docs, embedding_model) | |
| # Returns a FAISS wrapper vector store. Input is a list of strings. from_documents method used documents to Return VectorStore | |
| #data = dataset["text"] | |
| data = dataset["train"] | |
| print(data) | |
| d = 384 # vectors dimension | |
| m = 32 # hnsw parameter. Higher is more accurate but takes more time to index (default is 32, 128 should be ok) | |
| index = faiss.IndexHNSWFlat(d, m, faiss.METRIC_INNER_PRODUCT) | |
| data = data.add_faiss_index("embeddings", custom_index=index) | |
| # adds an index column that for the embeddings | |
| 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.""" | |
| # Provides context of how to answer the question | |
| print("check2") | |
| llm_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" | |
| tokenizer = AutoTokenizer.from_pretrained(llm_model) | |
| # pulling tokeinzer for text generation model | |
| 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 = 3 ): | |
| """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['text'][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 = 1 # number of retrieved documents | |
| scores , retrieved_documents = search(prompt, k) # get retrival scores and examples in dictionary format based on the promt passed | |
| formatted_prompt = format_prompt(prompt,retrieved_documents,k) # create a new prompt using the retrieved documents | |
| formatted_prompt = formatted_prompt[:400] # to avoid memory issue | |
| 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=300, | |
| eos_token_id=terminators, | |
| do_sample=True, | |
| temperature=0.6, | |
| 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= 512, | |
| do_sample=True, | |
| top_p=0.95, | |
| temperature=0.75, | |
| 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") | |
| TITLE = "AI Copilot for Diabetes Patients" | |
| DESCRIPTION = "" | |
| 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(share=True) |