DocChat_n_Talk / app.py
capradeepgujaran's picture
Update app.py
9d47d09 verified
raw
history blame
9.4 kB
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 sentence_transformers import SentenceTransformer, util
import logging
from openai_tts_tool import generate_audio_and_text
import tempfile
# Set up logging configuration
logging.basicConfig(level=logging.INFO, format='%(asctime)s | %(levelname)s | %(message)s')
# Initialize global variables
vector_index = None
query_log = []
sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
# Get available languages for OCR
try:
langs = os.popen('tesseract --list-langs').read().split('\n')[1:-1]
except:
langs = ['eng'] # Fallback to English if tesseract isn't properly configured
def create_temp_dir():
"""Create temporary directory if it doesn't exist"""
temp_dir = os.path.join(os.getcwd(), 'temp')
if not os.path.exists(temp_dir):
os.makedirs(temp_dir)
return temp_dir
def preprocess_image(image_path):
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_dir = create_temp_dir()
temp_filename = os.path.join(temp_dir, "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)
try:
os.remove(processed_image_path)
except:
pass
return text
def extract_text_from_pdf(pdf_path, lang='eng'):
text = ""
temp_dir = create_temp_dir()
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:
temp_image_path = os.path.join(temp_dir, f'temp_image_{page_num}.png')
image.save(temp_image_path, 'PNG')
text += extract_text_from_image(temp_image_path, lang=lang)
text += f"\n[OCR applied on page {page_num + 1}]\n"
try:
os.remove(temp_image_path)
except:
pass
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."
if not files:
return "No files uploaded."
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 {os.path.basename(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
except Exception as e:
return f"Error creating index: {str(e)}"
else:
return f"No valid documents were indexed. Errors: {'; '.join(error_messages)}"
def query_app(query, model_name, use_similarity_check, api_key):
global vector_index, query_log
if vector_index is None:
return "No documents indexed yet. Please upload documents first."
if not api_key:
return "Please provide a valid OpenAI API Key."
try:
llm = OpenAI(model=model_name, api_key=api_key)
response_synthesizer = get_response_synthesizer(llm=llm)
query_engine = vector_index.as_query_engine(llm=llm, response_synthesizer=response_synthesizer)
response = query_engine.query(query)
generated_response = response.response
return generated_response
except Exception as e:
logging.error(f"Error during query processing: {e}")
return f"Error during query processing: {str(e)}"
def create_gradio_interface():
with gr.Blocks(title="Document Processing and TTS App") as demo:
gr.Markdown("# πŸ“„ Document Processing, Text & Audio Generation App")
with gr.Tab("πŸ“€ Upload Documents"):
api_key_input = gr.Textbox(
label="Enter OpenAI API Key",
placeholder="Paste your OpenAI API Key here",
type="password"
)
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)
with gr.Tab("❓ Ask a Question"):
query_input = gr.Textbox(label="Enter your question")
model_dropdown = gr.Dropdown(
choices=["gpt-4-0125-preview", "gpt-3.5-turbo-0125"],
label="Select Model",
value="gpt-3.5-turbo-0125"
)
similarity_checkbox = gr.Checkbox(label="Use Similarity Check", value=False)
query_button = gr.Button("Ask")
answer_output = gr.Textbox(label="Answer", interactive=False)
with gr.Tab("πŸ—£οΈ Generate Audio and Text"):
text_input = gr.Textbox(label="Enter text for generation")
voice_type = gr.Dropdown(
choices=["alloy", "echo", "fable", "onyx", "nova", "shimmer"],
label="Voice Type",
value="alloy"
)
voice_speed = gr.Slider(
minimum=0.25,
maximum=4.0,
value=1.0,
label="Voice Speed"
)
language = gr.Dropdown(
choices=["en", "ar", "de", "hi", "es", "fr", "it", "ja", "ko", "pt"],
label="Language",
value="en"
)
output_option = gr.Radio(
choices=["audio", "summary_text", "both"],
label="Output Option",
value="both"
)
summary_length = gr.Slider(
minimum=50,
maximum=500,
value=100,
step=10,
label="Summary Length (words)"
)
additional_prompt = gr.Textbox(label="Additional Prompt (Optional)")
generate_button = gr.Button("Generate")
audio_output = gr.Audio(label="Generated Audio")
summary_output = gr.Textbox(label="Generated Summary Text")
# Wire up the components
upload_button.click(
fn=process_upload,
inputs=[api_key_input, file_upload, lang_dropdown],
outputs=[upload_status]
)
query_button.click(
fn=query_app,
inputs=[query_input, model_dropdown, similarity_checkbox, api_key_input],
outputs=[answer_output]
)
generate_button.click(
fn=generate_audio_and_text,
inputs=[
api_key_input, text_input, model_dropdown, voice_type,
voice_speed, language, output_option, summary_length,
additional_prompt
],
outputs=[audio_output, summary_output]
)
return demo
if __name__ == "__main__":
demo = create_gradio_interface()
demo.launch()
else:
demo = create_gradio_interface()