# -*- coding: utf-8 -*- """corrective RAG.ipynb Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/1x15B2MqaoYKN8vmvaDIMjhhFSSAELiDh """ CHROMA_DB_PATH = "chromadb" from langchain_community.document_loaders import PyPDFLoader from langchain_experimental.text_splitter import SemanticChunker from langchain_community.embeddings import GPT4AllEmbeddings import chromadb from transformers import AutoModelForCausalLM, AutoTokenizer import spacy from sentence_transformers.util import cos_sim import numpy as np from langchain_community.tools.tavily_search import TavilySearchResults from langchain_core.documents import Document from langchain_community.vectorstores import Chroma from langchain.schema import Document import gradio as gr import os from huggingface_hub import login import torch torch.set_default_device("cpu") hf_token = os.getenv("hftoken") login(token=hf_token) model_id = "mistralai/Mistral-7B-Instruct-v0.3" tokenizer = AutoTokenizer.from_pretrained(model_id,use_fast=True) model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="cpu",offload_folder="offload",offload_state_dict=True,) def extract_text_with_pypdf(pdf_path): loader = PyPDFLoader(pdf_path) pages = loader.load() text = "\n".join([page.page_content for page in pages]) return text def store_docs(): text = extract_text_with_pypdf("Build a Large Language Model.pdf")+extract_text_with_pypdf("Hands-On Large Language Model.pdf") embedding_model = GPT4AllEmbeddings() chunker = SemanticChunker(embedding_model) chunks = chunker.split_text(text) documents = [Document(page_content=chunk) for chunk in chunks] vector_store = Chroma.from_documents( documents=documents, embedding=embedding_model, persist_directory=CHROMA_DB_PATH ) vector_store.persist() def retrieve_docs(question:str)->list: embedding_model = GPT4AllEmbeddings() vector_store = Chroma(persist_directory=CHROMA_DB_PATH, embedding_function=embedding_model) results = vector_store.similarity_search(question, k=3) return [doc.page_content for doc in results] if results else [] def evaluate_docs(docs: list, question: str) -> list: results = [] for doc in docs: prompt = f""" Given a question, does the following document have exact information to answer the question? Question: {question} Document: {doc} Think step by step and classify the document as one of the following: - Correct (fully answers the question) - Incorrect (not relevant at all) - Ambiguous (partially relevant but unclear) Answer with one of these words only: Correct, Incorrect, or Ambiguous. """ inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate(**inputs, max_new_tokens=100) classification = tokenizer.decode(outputs[0], skip_special_tokens=True) classification = classification.split()[-1].strip() results.append((doc, classification)) return results def rewrite_question(question:str)->str: prompt=f""" Extract at most three keywords separated by comma from the following dialogues and questions as queries for the web search, including topic background within dialogues and main intent within questions. question: What is Henry Feilden’s occupation? query: Henry Feilden, occupation question: In what city was Billy Carlson born? query: city, Billy Carlson, born question: What is the religion of John Gwynn? query: religion of John Gwynn question: What sport does Kiribati men’s national basketball team play? query: sport, Kiribati men’s national basketball team play question: {question} query: """ inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate(**inputs, max_new_tokens=100) return tokenizer.decode(outputs[0], skip_special_tokens=True) def knowledge_refinement(doc:str, question:str)->str: nlp = spacy.load("en_core_web_sm") doc = nlp(doc) sentences = [sent.text for sent in doc.sents] embedding_model = GPT4AllEmbeddings() question_vector = np.array(embedding_model.embed_query(question)) sentence_vectors = [np.array(embedding_model.embed_query(sent)) for sent in sentences] similarities = [cos_sim(question_vector, doc_vector) for doc_vector in sentence_vectors] ranked_sentences = sorted(zip(sentences, similarities), key=lambda x: x[1], reverse=True) top_sentences = [s[0] for s in (ranked_sentences[:3] if len(ranked_sentences) > 3 else ranked_sentences)] return top_sentences def web_search(question:str)->list: question=rewrite_question(question) web_search_tool = TavilySearchResults(k=3) web_results = web_search_tool.invoke({"query": question}) return [d["content"] for d in web_results] def generation(documents:list, question:str)->str: prompt = f"""You are an assistant for question-answering tasks. Use the following documents to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise: Question: {question} Documents: {documents} Answer: """ inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate(**inputs, max_new_tokens=1000) return tokenizer.decode(outputs[0], skip_special_tokens=True) store_docs() def corrective_rag_pipeline(question): final_docs = [] result = evaluate_docs(retrieve_docs(question), question) correct_docs = [doc for doc, label in result if label == "Correct"] ambiguous_docs = [doc for doc, label in result if label == "Ambiguous"] if correct_docs: refined_docs = [knowledge_refinement(doc, question) for doc in correct_docs] final_docs.extend([item for sublist in refined_docs for item in sublist]) else: query = rewrite_question(question) web_docs = web_search(query) final_docs.extend(ambiguous_docs + [d.page_content for d in web_docs]) return generation(final_docs, question) with gr.Blocks() as demo: gr.Markdown("# **LLM Q&A Chatbot**") with gr.Row(): question_input = gr.Textbox(label="Ask a question", interactive=True) submit_button = gr.Button("Generate Answer") output_text = gr.Textbox(label="Answer", interactive=False) submit_button.click(corrective_rag_pipeline, inputs=question_input, outputs=output_text) demo.launch(share=True)