File size: 8,522 Bytes
8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c 3ddea46 8c44d5c 3ddea46 8c44d5c 3ddea46 8c44d5c 3ddea46 8c44d5c 3ddea46 8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c 60eae25 8c44d5c eb3a785 8c44d5c 60eae25 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 |
from typing import Any, List, Tuple
import gradio as gr
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import PyMuPDFLoader
import fitz
from PIL import Image
import os
import re
import openai
# MyApp class to handle the processes
class MyApp:
def __init__(self) -> None:
self.OPENAI_API_KEY: str = None # Initialize with None
self.chain = None
self.chat_history: list = []
self.documents = None
self.file_name = None
def set_api_key(self, api_key: str):
self.OPENAI_API_KEY = api_key
openai.api_key = api_key
def process_file(self, file) -> Image.Image:
loader = PyMuPDFLoader(file.name)
self.documents = loader.load()
self.file_name = os.path.basename(file.name)
doc = fitz.open(file.name)
page = doc[0]
pix = page.get_pixmap(dpi=150)
image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
return image
def build_chain(self, file) -> str:
embeddings = OpenAIEmbeddings(openai_api_key=self.OPENAI_API_KEY)
pdfsearch = Chroma.from_documents(
self.documents,
embeddings,
collection_name=self.file_name,
)
self.chain = ConversationalRetrievalChain.from_llm(
ChatOpenAI(temperature=0.0, openai_api_key=self.OPENAI_API_KEY),
retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),
return_source_documents=True,
)
return "Vector database built successfully!"
# Function to add text to chat history
def add_text(history: List[Tuple[str, str]], text: str) -> List[Tuple[str, str]]:
if not text:
raise gr.Error("Enter text")
history.append((text, ""))
return history
# Function to get response from the model
def get_response(history, query):
if app.chain is None:
raise gr.Error("The chain has not been built yet. Please ensure the vector database is built before querying.")
try:
result = app.chain.invoke(
{"question": query, "chat_history": app.chat_history}
)
app.chat_history.append((query, result["answer"]))
source_docs = result["source_documents"]
source_texts = []
for doc in source_docs:
source_texts.append(f"Page {doc.metadata['page'] + 1}: {doc.page_content}")
source_texts_str = "\n\n".join(source_texts)
history[-1] = (history[-1][0], result["answer"])
return history, source_texts_str
except Exception as e:
app.chat_history.append((query, "I have no information about it. Feed me knowledge, please!"))
return history, f"I have no information about it. Feed me knowledge, please! Error: {str(e)}"
# Function to get response for the current RAG tab
def get_response_current(history, query):
if app.chain is None:
raise gr.Error("The chain has not been built yet. Please ensure the vector database is built before querying.")
try:
result = app.chain.invoke(
{"question": query, "chat_history": app.chat_history}
)
app.chat_history.append((query, result["answer"]))
source_docs = result["source_documents"]
source_texts = []
for doc in source_docs:
source_texts.append(f"Page {doc.metadata['page'] + 1}: {doc.page_content}")
source_texts_str = "\n\n".join(source_texts)
history[-1] = (history[-1][0], result["answer"])
return history, source_texts_str
except Exception as e:
app.chat_history.append((query, "I have no information about it. Feed me knowledge, please!"))
return history, f"I have no information about it. Feed me knowledge, please! Error: {str(e)}"
# Function to render file
def render_file(file) -> Image.Image:
doc = fitz.open(file.name)
page = doc[0]
pix = page.get_pixmap(dpi=150)
image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
return image
# Function to purge chat and render first page of PDF
def purge_chat_and_render_first(file) -> Tuple[Image.Image, list]:
app.chat_history = []
doc = fitz.open(file.name)
page = doc[0]
pix = page.get_pixmap(dpi=150)
image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
return image, []
# Function to refresh chat
def refresh_chat():
app.chat_history = []
return []
app = MyApp()
# Function to set API key
def set_api_key(api_key):
app.set_api_key(api_key)
# Pre-process the saved PDF file after setting the API key
saved_file_path = "track_training.pdf"
with open(saved_file_path, 'rb') as saved_file:
app.process_file(saved_file)
app.build_chain(saved_file)
return f"API Key set to {api_key[:4]}...{api_key[-4:]} and vector database built successfully!"
# Gradio interface
with gr.Blocks() as demo:
api_key_input = gr.Textbox(label="OpenAI API Key", type="password", placeholder="Enter your OpenAI API Key")
api_key_btn = gr.Button("Set API Key")
api_key_status = gr.Textbox(value="API Key status", interactive=False)
api_key_btn.click(
fn=set_api_key,
inputs=[api_key_input],
outputs=[api_key_status]
)
with gr.Tab("Inst RAG"):
with gr.Tab("Upload PDF"):
btn = gr.UploadButton("📁 Upload a PDF", file_types=[".pdf"])
show_img = gr.Image(label="Uploaded PDF")
btn.upload(
fn=purge_chat_and_render_first,
inputs=[btn],
outputs=[show_img],
)
with gr.Tab("Process PDF"):
process_btn = gr.Button("Process PDF")
show_img_processed = gr.Image(label="Processed PDF")
process_status = gr.Textbox(label="Processing Status", interactive=False)
process_btn.click(
fn=lambda file: (app.process_file(file), "Processing complete!"),
inputs=[btn],
outputs=[show_img_processed, process_status],
)
with gr.Tab("Build Vector Database"):
build_vector_btn = gr.Button("Build Vector Database")
status_text = gr.Textbox(label="Status", value="", interactive=False)
build_vector_btn.click(
fn=app.build_chain,
inputs=[btn],
outputs=[status_text],
)
with gr.Tab("Chat"):
chatbot = gr.Chatbot(elem_id="chatbot")
txt = gr.Textbox(
show_label=False,
placeholder="Enter text and press submit",
scale=2
)
submit_btn = gr.Button("Submit", scale=1)
refresh_btn = gr.Button("Refresh Chat", scale=1)
source_texts_output = gr.Textbox(label="Source Texts", interactive=False)
submit_btn.click(
fn=add_text,
inputs=[chatbot, txt],
outputs=[chatbot],
queue=False,
).success(
fn=get_response, inputs=[chatbot, txt], outputs=[chatbot, source_texts_output]
)
refresh_btn.click(
fn=refresh_chat,
inputs=[],
outputs=[chatbot],
)
with gr.Tab("Current RAG"):
with gr.Column():
chatbot_current = gr.Chatbot(elem_id="chatbot_current")
txt_current = gr.Textbox(
show_label=False,
placeholder="Enter text and press submit",
scale=2
)
submit_btn_current = gr.Button("Submit", scale=1)
refresh_btn_current = gr.Button("Refresh Chat", scale=1)
source_texts_output_current = gr.Textbox(label="Source Texts", interactive=False)
submit_btn_current.click(
fn=add_text,
inputs=[chatbot_current, txt_current],
outputs=[chatbot_current],
queue=False,
).success(
fn=get_response_current, inputs=[chatbot_current, txt_current], outputs=[chatbot_current, source_texts_output_current]
)
refresh_btn_current.click(
fn=refresh_chat,
inputs=[],
outputs=[chatbot_current],
)
demo.queue()
demo.launch()
|