app.py

# -*- 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)