|
import os |
|
import cv2 |
|
import numpy as np |
|
from PIL import Image |
|
import pytesseract |
|
import gradio as gr |
|
from pdf2image import convert_from_path |
|
import PyPDF2 |
|
from llama_index.core import VectorStoreIndex, Document |
|
from llama_index.embeddings.openai import OpenAIEmbedding |
|
from llama_index.llms.openai import OpenAI |
|
from llama_index.core import get_response_synthesizer |
|
from dotenv import load_dotenv |
|
from sentence_transformers import SentenceTransformer, util |
|
import logging |
|
|
|
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)s | %(message)s') |
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
vector_index = None |
|
query_log = [] |
|
sentence_model = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
|
langs = os.popen('tesseract --list-langs').read().split('\n')[1:-1] |
|
|
|
def preprocess_image(image_path): |
|
""" |
|
Pre-process the image to improve OCR results. |
|
- Convert to grayscale |
|
- Apply thresholding to improve contrast |
|
- Apply denoising if needed |
|
""" |
|
img = cv2.imread(image_path) |
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
|
gray = cv2.equalizeHist(gray) |
|
gray = cv2.GaussianBlur(gray, (5, 5), 0) |
|
processed_image = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, |
|
cv2.THRESH_BINARY, 11, 2) |
|
temp_filename = "processed_image.png" |
|
cv2.imwrite(temp_filename, processed_image) |
|
return temp_filename |
|
|
|
def extract_text_from_image(image_path, lang='eng'): |
|
processed_image_path = preprocess_image(image_path) |
|
text = pytesseract.image_to_string(Image.open(processed_image_path), lang=lang) |
|
return text |
|
|
|
def extract_text_from_pdf(pdf_path, lang='eng'): |
|
text = "" |
|
try: |
|
with open(pdf_path, 'rb') as file: |
|
pdf_reader = PyPDF2.PdfReader(file) |
|
for page_num in range(len(pdf_reader.pages)): |
|
page = pdf_reader.pages[page_num] |
|
page_text = page.extract_text() |
|
if page_text.strip(): |
|
text += page_text |
|
else: |
|
images = convert_from_path(pdf_path, first_page=page_num + 1, last_page=page_num + 1) |
|
for image in images: |
|
image.save('temp_image.png', 'PNG') |
|
text += extract_text_from_image('temp_image.png', lang=lang) |
|
text += f"\n[OCR applied on page {page_num + 1}]\n" |
|
except Exception as e: |
|
return f"Error processing PDF: {str(e)}" |
|
return text |
|
|
|
def extract_text(file_path, lang='eng'): |
|
file_ext = file_path.lower().split('.')[-1] |
|
if file_ext in ['pdf']: |
|
return extract_text_from_pdf(file_path, lang) |
|
elif file_ext in ['png', 'jpg', 'jpeg']: |
|
return extract_text_from_image(file_path, lang) |
|
else: |
|
return f"Unsupported file type: {file_ext}" |
|
|
|
def process_upload(api_key, files, lang): |
|
global vector_index |
|
|
|
if not api_key: |
|
return "Please provide a valid OpenAI API Key.", None |
|
|
|
if not files: |
|
return "No files uploaded.", None |
|
|
|
documents = [] |
|
error_messages = [] |
|
image_heavy_docs = [] |
|
|
|
for file_path in files: |
|
try: |
|
text = extract_text(file_path, lang) |
|
if "This document consists of" in text and "page(s) of images" in text: |
|
image_heavy_docs.append(os.path.basename(file_path)) |
|
documents.append(Document(text=text)) |
|
except Exception as e: |
|
error_message = f"Error processing file {file_path}: {str(e)}" |
|
logging.error(error_message) |
|
error_messages.append(error_message) |
|
|
|
if documents: |
|
try: |
|
embed_model = OpenAIEmbedding(model="text-embedding-3-large", api_key=api_key) |
|
vector_index = VectorStoreIndex.from_documents(documents, embed_model=embed_model) |
|
|
|
success_message = f"Successfully indexed {len(documents)} files." |
|
if image_heavy_docs: |
|
success_message += f"\nNote: The following documents consist mainly of images and may require manual review: {', '.join(image_heavy_docs)}" |
|
if error_messages: |
|
success_message += f"\nErrors: {'; '.join(error_messages)}" |
|
|
|
return success_message, vector_index |
|
except Exception as e: |
|
return f"Error creating index: {str(e)}", None |
|
else: |
|
return f"No valid documents were indexed. Errors: {'; '.join(error_messages)}", None |
|
|
|
def calculate_similarity(response, ground_truth): |
|
response_embedding = sentence_model.encode(response, convert_to_tensor=True) |
|
truth_embedding = sentence_model.encode(ground_truth, convert_to_tensor=True) |
|
|
|
response_embedding = response_embedding / np.linalg.norm(response_embedding) |
|
truth_embedding = truth_embedding / np.linalg.norm(truth_embedding) |
|
|
|
similarity = np.dot(response_embedding, truth_embedding) |
|
similarity_percentage = (similarity + 1) / 2 * 100 |
|
|
|
return similarity_percentage |
|
|
|
def query_app(query, model_name, use_similarity_check, openai_api_key): |
|
global vector_index, query_log |
|
|
|
if vector_index is None: |
|
logging.error("No documents indexed yet. Please upload documents first.") |
|
return "No documents indexed yet. Please upload documents first.", None |
|
|
|
if not openai_api_key: |
|
logging.error("No OpenAI API Key provided.") |
|
return "Please provide a valid OpenAI API Key.", None |
|
|
|
try: |
|
llm = OpenAI(model=model_name, api_key=openai_api_key) |
|
except Exception as e: |
|
logging.error(f"Error initializing the OpenAI model: {e}") |
|
return f"Error initializing the OpenAI model: {e}", None |
|
|
|
response_synthesizer = get_response_synthesizer(llm=llm) |
|
query_engine = vector_index.as_query_engine(llm=llm, response_synthesizer=response_synthesizer) |
|
|
|
try: |
|
response = query_engine.query(query) |
|
except Exception as e: |
|
logging.error(f"Error during query processing: {e}") |
|
return f"Error during query processing: {e}", None |
|
|
|
generated_response = response.response |
|
query_log.append({ |
|
"query_id": str(len(query_log) + 1), |
|
"query": query, |
|
"gt_answer": "Placeholder ground truth answer", |
|
"response": generated_response, |
|
"retrieved_context": [{"text": doc.text} for doc in response.source_nodes] |
|
}) |
|
|
|
metrics = {} |
|
|
|
if use_similarity_check: |
|
try: |
|
logging.info("Similarity check is enabled. Calculating similarity.") |
|
similarity = calculate_similarity(generated_response, "Placeholder ground truth answer") |
|
metrics['similarity'] = similarity |
|
logging.info(f"Similarity calculated: {similarity}") |
|
except Exception as e: |
|
logging.error(f"Error during similarity calculation: {e}") |
|
metrics['error'] = f"Error during similarity calculation: {e}" |
|
|
|
return generated_response, metrics if use_similarity_check else None |
|
|
|
def main(): |
|
with gr.Blocks(title="Document Processing App") as demo: |
|
gr.Markdown("# 📄 Document Processing and Querying App") |
|
|
|
with gr.Tab("📤 Upload Documents"): |
|
gr.Markdown("### Enter your OpenAI API Key and Upload PDF, DOCX, or TXT files to index") |
|
|
|
api_key_input = gr.Textbox(label="Enter OpenAI API Key", placeholder="Paste your OpenAI API Key here") |
|
|
|
with gr.Row(): |
|
file_upload = gr.File(label="Upload Files", file_count="multiple", type="filepath") |
|
lang_dropdown = gr.Dropdown(choices=langs, label="Select OCR Language", value='eng') |
|
upload_button = gr.Button("Upload and Index") |
|
upload_status = gr.Textbox(label="Status", interactive=False) |
|
|
|
upload_button.click( |
|
fn=process_upload, |
|
inputs=[api_key_input, file_upload, lang_dropdown], |
|
outputs=[upload_status] |
|
) |
|
|
|
with gr.Tab("❓ Ask a Question"): |
|
gr.Markdown("### Query the indexed documents") |
|
with gr.Column(): |
|
query_input = gr.Textbox(label="Enter your question", placeholder="Type your question here...") |
|
model_dropdown = gr.Dropdown( |
|
choices=["gpt-4o", "gpt-4o-mini"], |
|
value="gpt-4o", |
|
label="Select Model" |
|
) |
|
similarity_checkbox = gr.Checkbox(label="Use Similarity Check", value=False) |
|
query_button = gr.Button("Ask") |
|
with gr.Column(): |
|
answer_output = gr.Textbox(label="Answer", interactive=False) |
|
metrics_output = gr.JSON(label="Metrics") |
|
|
|
query_button.click( |
|
fn=query_app, |
|
inputs=[query_input, model_dropdown, similarity_checkbox, api_key_input], |
|
outputs=[answer_output, metrics_output] |
|
) |
|
|
|
gr.Markdown(""" |
|
--- |
|
**Note:** Ensure you upload documents before attempting to query. Enter a valid OpenAI API Key to interact with the models. |
|
""") |
|
|
|
demo.launch() |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
