| # import torch | |
| # import transformers | |
| # from transformers import RagRetriever, RagSequenceForGeneration, AutoTokenizer, AutoModelForCausalLM | |
| # import gradio as gr | |
| # device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| # dataset_path = "./5k_index_data/my_knowledge_dataset" | |
| # index_path = "./5k_index_data/my_knowledge_dataset_hnsw_index.faiss" | |
| # tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq") | |
| # retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom", | |
| # passages_path = dataset_path, | |
| # index_path = index_path, | |
| # n_docs = 5) | |
| # rag_model = RagSequenceForGeneration.from_pretrained('facebook/rag-sequence-nq', retriever=retriever) | |
| # rag_model.retriever.init_retrieval() | |
| # rag_model.to(device) | |
| # model = AutoModelForCausalLM.from_pretrained('HuggingFaceH4/zephyr-7b-beta', | |
| # device_map = 'auto', | |
| # torch_dtype = torch.bfloat16, | |
| # ) | |
| # def strip_title(title): | |
| # if title.startswith('"'): | |
| # title = title[1:] | |
| # if title.endswith('"'): | |
| # title = title[:-1] | |
| # return title | |
| # # getting the correct format to input in gemma model | |
| # def input_format(query, context): | |
| # sys_instruction = f'Context:\n {context} \n Given the following information, generate answer to the question. Provide links in the answer from the information to increase credebility.' | |
| # message = f'Question: {query}' | |
| # return f'<bos><start_of_turn>\n{sys_instruction}' + f' {message}<end_of_turn>\n' | |
| # # retrieving and generating answer in one call | |
| # def retrieved_info(query, rag_model = rag_model, generating_model = model): | |
| # # Tokenize Query | |
| # retriever_input_ids = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus( | |
| # [query], | |
| # return_tensors = 'pt', | |
| # padding = True, | |
| # truncation = True, | |
| # )['input_ids'].to(device) | |
| # # Retrieve Documents | |
| # question_encoder_output = rag_model.rag.question_encoder(retriever_input_ids) | |
| # question_encoder_pool_output = question_encoder_output[0] | |
| # result = rag_model.retriever( | |
| # retriever_input_ids, | |
| # question_encoder_pool_output.cpu().detach().to(torch.float32).numpy(), | |
| # prefix = rag_model.rag.generator.config.prefix, | |
| # n_docs = rag_model.config.n_docs, | |
| # return_tensors = 'pt', | |
| # ) | |
| # # Preparing query and retrieved docs for model | |
| # all_docs = rag_model.retriever.index.get_doc_dicts(result.doc_ids) | |
| # retrieved_context = [] | |
| # for docs in all_docs: | |
| # titles = [strip_title(title) for title in docs['title']] | |
| # texts = docs['text'] | |
| # for title, text in zip(titles, texts): | |
| # retrieved_context.append(f'{title}: {text}') | |
| # generation_model_input = input_format(query, retrieved_context) | |
| # # Generating answer using gemma model | |
| # tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta") | |
| # input_ids = tokenizer(generation_model_input, return_tensors='pt').to(device) | |
| # output = generating_model.generate(input_ids, max_new_tokens = 256) | |
| # return tokenizer.decode(output[0]) | |
| # def respond( | |
| # message, | |
| # history: list[tuple[str, str]], | |
| # system_message, | |
| # max_tokens , | |
| # temperature, | |
| # top_p, | |
| # ): | |
| # if message: # If there's a user query | |
| # response = retrieved_info(message) # Get the answer from your local FAISS and Q&A model | |
| # return response | |
| # # In case no message, return an empty string | |
| # return "" | |
| # """ | |
| # For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface | |
| # """ | |
| # # Custom title and description | |
| # title = "🧠 Welcome to Your AI Knowledge Assistant" | |
| # description = """ | |
| # Hi!! I am your loyal assistant. My functionality is based on the RAG model. I retrieve relevant information and provide answers based on that. Ask me any questions, and let me assist you. | |
| # My capabilities are limited because I am still in the development phase. I will do my best to assist you. SOOO LET'S BEGGINNNN...... | |
| # """ | |
| # demo = gr.ChatInterface( | |
| # respond, | |
| # type = 'messages', | |
| # additional_inputs=[ | |
| # gr.Textbox(value="You are a helpful and friendly assistant.", label="System message"), | |
| # gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="Max new tokens"), | |
| # gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), | |
| # gr.Slider( | |
| # minimum=0.1, | |
| # maximum=1.0, | |
| # value=0.95, | |
| # step=0.05, | |
| # label="Top-p (nucleus sampling)", | |
| # ), | |
| # ], | |
| # title=title, | |
| # description=description, | |
| # textbox=gr.Textbox(placeholder=["'What is the future of AI?' or 'App Development'"]), | |
| # examples=[["✨Future of AI"], ["📱App Development"]], | |
| # example_icons=["🤖", "📱"], | |
| # theme="compact", | |
| # submit_btn = True, | |
| # ) | |
| # if __name__ == "__main__": | |
| # demo.launch(share = True ) | |
| import torch | |
| import transformers | |
| from transformers import RagRetriever, RagSequenceForGeneration, AutoTokenizer, AutoModelForCausalLM | |
| import gradio as gr | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| dataset_path = "./5k_index_data/my_knowledge_dataset" | |
| index_path = "./5k_index_data/my_knowledge_dataset_hnsw_index.faiss" | |
| tokenizer = AutoTokenizer.from_pretrained("facebook/rag-sequence-nq") | |
| retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom", | |
| passages_path = dataset_path, | |
| index_path = index_path, | |
| n_docs = 5) | |
| rag_model = RagSequenceForGeneration.from_pretrained('facebook/rag-sequence-nq', retriever=retriever) | |
| rag_model.retriever.init_retrieval() | |
| rag_model.to(device) | |
| model = AutoModelForCausalLM.from_pretrained('HuggingFaceH4/zephyr-7b-beta', | |
| device_map = 'auto', | |
| torch_dtype = torch.bfloat16, | |
| ) | |
| def strip_title(title): | |
| if title.startswith('"'): | |
| title = title[1:] | |
| if title.endswith('"'): | |
| title = title[:-1] | |
| return title | |
| # getting the correct format to input in gemma model | |
| def input_format(query, context): | |
| # sys_instruction = f'Context:\n {context} \n Given the following information, generate answer to the question. Provide links in the answer from the information to increase credebility.' | |
| # message = f'Question: {query}' | |
| # return f'<bos><start_of_turn>\n{sys_instruction}' + f' {message}<end_of_turn>\n' | |
| return [ | |
| { | |
| "role": "system", "content": f'Context:\n {context} \n Given the following information, generate answer to the question. Provide links in the answer from the information to increase credebility.' }, | |
| { | |
| "role": "user", "content": f"{query}"}, | |
| ] | |
| # retrieving and generating answer in one call | |
| def retrieved_info(query, rag_model = rag_model, generating_model = model): | |
| # Tokenize Query | |
| retriever_input_ids = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus( | |
| [query], | |
| return_tensors = 'pt', | |
| padding = True, | |
| truncation = True, | |
| )['input_ids'].to(device) | |
| # Retrieve Documents | |
| question_encoder_output = rag_model.rag.question_encoder(retriever_input_ids) | |
| question_encoder_pool_output = question_encoder_output[0] | |
| result = rag_model.retriever( | |
| retriever_input_ids, | |
| question_encoder_pool_output.cpu().detach().to(torch.float32).numpy(), | |
| prefix = rag_model.rag.generator.config.prefix, | |
| n_docs = rag_model.config.n_docs, | |
| return_tensors = 'pt', | |
| ) | |
| # Preparing query and retrieved docs for model | |
| all_docs = rag_model.retriever.index.get_doc_dicts(result.doc_ids) | |
| retrieved_context = [] | |
| for docs in all_docs: | |
| titles = [strip_title(title) for title in docs['title']] | |
| texts = docs['text'] | |
| for title, text in zip(titles, texts): | |
| retrieved_context.append(f'{title}: {text}') | |
| print(retrieved_context) | |
| generation_model_input = input_format(query, retrieved_context[0]) | |
| # Generating answer using gemma model | |
| tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta") | |
| input_ids = tokenizer(generation_model_input, return_tensors='pt')['input_ids'].to(device) | |
| output = generating_model.generate(input_ids, max_new_tokens = 256) | |
| return tokenizer.decode(output[0]) | |
| def respond( | |
| message, | |
| history: list[tuple[str, str]], | |
| system_message, | |
| max_tokens , | |
| temperature, | |
| top_p, | |
| ): | |
| if message: # If there's a user query | |
| response = retrieved_info(message) # Get the answer from your local FAISS and Q&A model | |
| return response | |
| # In case no message, return an empty string | |
| return "" | |
| """ | |
| For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface | |
| """ | |
| # Custom title and description | |
| title = "🧠 Welcome to Your AI Knowledge Assistant" | |
| description = """ | |
| Hi!! I am your loyal assistant. My functionality is based on the RAG model. I retrieve relevant information and provide answers based on that. Ask me any questions, and let me assist you. | |
| My capabilities are limited because I am still in the development phase. I will do my best to assist you. SOOO LET'S BEGGINNNN...... | |
| """ | |
| demo = gr.ChatInterface( | |
| respond, | |
| type = 'messages', | |
| additional_inputs=[ | |
| gr.Textbox(value="You are a helpful and friendly assistant.", label="System message"), | |
| gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="Max new tokens"), | |
| gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), | |
| gr.Slider( | |
| minimum=0.1, | |
| maximum=1.0, | |
| value=0.95, | |
| step=0.05, | |
| label="Top-p (nucleus sampling)", | |
| ), | |
| ], | |
| title=title, | |
| description=description, | |
| textbox=gr.Textbox(placeholder=["'What is the future of AI?' or 'App Development'"]), | |
| examples=[["✨Future of AI"], ["📱App Development"]], | |
| #example_icons=["🤖", "📱"], | |
| theme="compact", | |
| submit_btn = True, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(share = True, | |
| show_error = True) | |