Spaces:
Sleeping
Sleeping
File size: 5,321 Bytes
9406eac 6f14fd9 9406eac e657d8c 9406eac e657d8c 9406eac |
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 |
import os
import tempfile
import time
import re
import logging
from datetime import datetime
import gradio as gr
import google.generativeai as genai
from PyPDF2 import PdfReader
from tika import parser
# Configure logging
tmp_log = "pdf_processor_log.txt"
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s',
handlers=[
logging.StreamHandler(),
logging.FileHandler(tmp_log)
]
)
logger = logging.getLogger("pdf_processor")
# Attempt to import Unstructured.io partitioning
try:
from unstructured.partition.pdf import partition_pdf
UNSTRUCTURED_AVAILABLE = True
except ImportError:
UNSTRUCTURED_AVAILABLE = False
logger.warning("unstructured.partition.pdf not available; skipping that extraction method")
# Load API key from environment
API_KEY = os.getenv("GOOGLE_API_KEY", None)
if not API_KEY:
logger.warning("GOOGLE_API_KEY not set in environment.")
else:
genai.configure(api_key=API_KEY)
# Globals to store state
EXTRACTED_TEXT = ""
PDF_SECTIONS = []
EXTRACTION_METHOD = ""
# --- Extraction Functions ---
def extract_text_with_unstructured(pdf_path):
logger.info("Extracting via Unstructured.io...")
elements = partition_pdf(filename=pdf_path, extract_images_in_pdf=False)
sections, current = [], {"title":"Introduction","content":""}
for e in elements:
if hasattr(e, "text") and (t := e.text.strip()):
if len(t)<80 and (t.isupper() or t.endswith(':') or re.match(r'^[0-9]+\.?\s+', t)):
if current["content"]: sections.append(current)
current = {"title":t, "content":""}
else:
current["content"] += t + "\n\n"
if current["content"]: sections.append(current)
return sections
def extract_text_with_pypdf(pdf_path):
logger.info("Extracting via PyPDF2...")
reader = PdfReader(pdf_path)
full = ""
for i,p in enumerate(reader.pages,1):
if (txt := p.extract_text()): full += f"\n\n--- Page {i} ---\n\n{txt}"
parts = re.split(r"\n\s*([A-Z][A-Z\s]+:?|[0-9]+\.\s+[A-Z].*?)\s*\n", full)
if len(parts)>1:
return [{"title":parts[i].strip(),"content":parts[i+1].strip()} for i in range(1,len(parts),2)]
# fallback to single section
return [{"title":"Document","content":full}]
def extract_text_with_tika(pdf_path):
logger.info("Extracting via Tika...")
parsed = parser.from_file(pdf_path)
lines = parsed.get("content","").split("\n")
sections, current = [], {"title":"Introduction","content":""}
for ln in lines:
ln = ln.strip()
if not ln: continue
if len(ln)<80 and (ln.isupper() or ln.endswith(':') or re.match(r'^[0-9]+\.?\s+[A-Z]', ln)):
if current["content"]: sections.append(current)
current = {"title":ln, "content":""}
else:
current["content"] += ln + "\n\n"
if current["content"]: sections.append(current)
return sections
# --- Gemini API calls ---
def generate_greg_brockman_summary(content):
model = genai.GenerativeModel('gemini-1.5-pro')
prompt = f"""
You are an expert document analyst...
{content}
"""
try:
resp = model.generate_content(prompt)
return resp.text, None
except Exception as e:
logger.error(e)
return None, str(e)
def answer_question_about_pdf(content, question):
model = genai.GenerativeModel('gemini-1.5-pro')
prompt = f"""
You are a precise document analysis assistant...
DOCUMENT CONTENT:
{content}
QUESTION: {question}
"""
try:
resp = model.generate_content(prompt)
return resp.text, None
except Exception as e:
logger.error(e)
return None, str(e)
# --- Processing & Q&A ---
def process_pdf(pdf_file, progress=gr.Progress()):
global EXTRACTED_TEXT, PDF_SECTIONS, EXTRACTION_METHOD
if not API_KEY:
return None, None, "❌ Set GOOGLE_API_KEY in settings.", ""
if pdf_file is None:
return None, None, "❌ No file uploaded.", ""
tmp = tempfile.gettempdir()
path = os.path.join(tmp, pdf_file.name)
with open(path, 'wb') as f: f.write(pdf_file.read())
methods = []
if UNSTRUCTURED_AVAILABLE:
methods.append(("unstructured", extract_text_with_unstructured))
methods.extend([
("pypdf", extract_text_with_pypdf),
("tika", extract_text_with_tika)
])
with gr.Tab("Ask Questions"):
question = gr.Textbox(label="Question", lines=2)
ask_btn = gr.Button("Ask")
answer = gr.Textbox(label="Answer", lines=10)
ask_btn.click(ask_question, inputs=[question], outputs=[answer])
with gr.Tab("System Log"):
refresh = gr.Button("Refresh Log")
syslog = gr.Textbox(label="System Log", lines=15)
refresh.click(view_log, inputs=None, outputs=[syslog])
with gr.Row():
save_sum_btn = gr.Button("Save Summary")
save_sum_status = gr.Markdown("")
save_sum_btn.click(save_summary, inputs=[summary_out], outputs=[save_sum_status])
with gr.Row():
save_qa_btn = gr.Button("Save Q&A")
save_qa_status = gr.Markdown("")
save_qa_btn.click(save_qa, inputs=[question, answer], outputs=[save_qa_status])
if __name__ == "__main__":
# For Hugging Face Spaces, set `server_name="0.0.0.0"` if needed
app.launch() |