File size: 9,217 Bytes
7e2fea5 95bd0c2 9019fe4 7e2fea5 ffbb3a2 7e2fea5 95bd0c2 | 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 | import html
import logging
import os
import threading
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
import spaces
import gradio as gr
from lexquery.config import GEMINI_MODEL, LANGUAGES
from lexquery.models import LocalModels
from lexquery.pipeline import Engine
from lexquery.schema import Result
logging.basicConfig(level=logging.INFO, format='%(levelname)s %(name)s: %(message)s')
@lru_cache(maxsize=1)
def load_models():
return LocalModels()
@dataclass
class Session:
engine: Engine = field(default_factory=lambda: Engine(load_models()))
chat: list = field(default_factory=list)
lock: object = field(default_factory=threading.RLock)
_sessions = {}
_sessions_lock = threading.RLock()
def session(request):
if not request.session_hash:
raise gr.Error('Open the app in a browser session first.')
with _sessions_lock:
if request.session_hash not in _sessions:
_sessions[request.session_hash] = Session()
return _sessions[request.session_hash]
def cleanup_id(session_id):
with _sessions_lock:
_sessions.pop(session_id, None)
def unload(request: gr.Request):
cleanup_id(request.session_hash)
def initialize(request: gr.Request):
session(request)
return request.session_hash
def safe_error(exc):
text = str(exc)
for name in ('GOOGLE_API_KEY', 'GEMINI_API_KEY'):
key = os.getenv(name)
if key:
text = text.replace(key, '[key hidden]')
return text
def read_uploads(files, request: gr.Request, progress=gr.Progress()):
item = session(request)
with item.lock:
keep, messages = set(), []
old_ids = set(item.engine.documents)
for i, file in enumerate(files or []):
path = Path(file)
progress(i / max(1, len(files)), desc='Reading ' + path.name)
try:
doc = item.engine.ingest(path.name, path.read_bytes())
keep.add(doc.id)
messages.append(f'**{html.escape(doc.name)}** · {len(doc.pages)} pages · ready')
uncertain = [str(p.number) for p in doc.pages if p.warnings]
if uncertain:
messages.append('Extraction notes on pages ' + ', '.join(uncertain) + '; see Extracted text.')
except Exception as exc:
logging.exception('PDF extraction failed')
messages.append(f'**{html.escape(path.name)}** · {html.escape(safe_error(exc))}')
for doc_id in set(item.engine.documents) - keep:
del item.engine.documents[doc_id]
if old_ids != keep:
item.engine.invalidate()
item.engine.history.clear()
item.chat.clear()
choices = [(d.name, d.id) for d in item.engine.documents.values()]
ids = [i for _, i in choices]
progress(1, desc='Finished')
return ('\n\n'.join(messages) or 'Upload a PDF to begin.',
gr.Dropdown(choices=choices, value=ids, multiselect=True),
gr.Dropdown(choices=choices, value=ids[0] if ids else None),
[], '', '')
def extracted_text(doc_id, request: gr.Request):
item = session(request)
with item.lock:
doc = item.engine.documents.get(doc_id)
if not doc:
return ''
return '\n\n'.join(f'PAGE {p.number} · {p.language}\n{p.original}\n' +
('\nNotes: ' + '; '.join(p.warnings) if p.warnings else '') for p in doc.pages)
def result_markdown(result):
paragraphs = [html.escape(text) for text in result.displayed] if result.status in (
'checked', 'conversational') else []
for note in result.notes:
paragraphs.append(html.escape(safe_error(note)))
return '\n\n'.join(paragraphs) or 'I could not find a supported answer in the selected documents.'
def ask(question, ids, language, request: gr.Request):
item = session(request)
with item.lock:
if not question.strip():
return item.chat, ''
result = item.engine.present(item.engine.ask(question, ids or []), language)
item.chat.extend([{'role': 'user', 'content': question},
{'role': 'assistant', 'content': result_markdown(result)}])
return item.chat, ''
@spaces.GPU
def ask(question, ids, language, request: gr.Request):
item = session(request)
with item.lock:
if not question.strip():
return item.chat, ''
result = item.engine.present(item.engine.ask(question, ids or []), language)
item.chat.extend([{'role': 'user', 'content': question},
{'role': 'assistant', 'content': result_markdown(result)}])
return item.chat, ''
@spaces.GPU
def summarize(ids, language, request: gr.Request, progress=gr.Progress()):
item = session(request)
with item.lock:
if not ids:
return 'Upload and select a PDF first.'
sections = []
def summarize(ids, language, request: gr.Request, progress=gr.Progress()):
item = session(request)
with item.lock:
if not ids:
return 'Upload and select a PDF first.'
sections = []
for i, doc_id in enumerate(ids):
doc = item.engine.documents.get(doc_id)
if doc is None:
continue
progress(i / len(ids), desc='Summarizing ' + doc.name)
try:
report = item.engine.report(doc_id)
result = item.engine.present(Result.model_validate(report['summary']), language)
section = '### ' + html.escape(doc.name) + '\n\n' + result_markdown(result)
fields = []
for name, facts in report['fields'].items():
if facts:
fields.append('**' + name.replace('_', ' ').title() + ':** ' +
'; '.join(html.escape(f['value']) for f in facts))
if fields:
section += '\n\n' + '\n\n'.join(fields)
for note in report.get('field_notes', []):
section += '\n\n' + html.escape(safe_error(note))
sections.append(section)
except Exception as exc:
logging.exception('Report failed')
sections.append('### ' + html.escape(doc.name) + '\n\n' + html.escape(safe_error(exc)))
progress(1, desc='Finished')
return '\n\n---\n\n'.join(sections)
def clear_chat(request: gr.Request):
item = session(request)
with item.lock:
item.engine.history.clear()
item.chat.clear()
return []
def build_app():
with gr.Blocks(title='LexQuery', delete_cache=(3600, 3600)) as demo:
lifecycle = gr.State(None, delete_callback=cleanup_id)
gr.Markdown('# LexQuery\nUpload your PDFs. Ask questions. Get clear summaries.')
with gr.Row():
with gr.Column(scale=1, min_width=280):
files = gr.File(label='Documents', file_types=['.pdf'], file_count='multiple', type='filepath')
gr.Markdown('Scanned and text PDFs · English and 22 Indian languages')
status = gr.Markdown('Upload a PDF to begin.')
selected = gr.Dropdown(choices=[], multiselect=True, label='Use these documents')
language = gr.Dropdown(LANGUAGES, value='English', label='Answer language')
with gr.Column(scale=3):
with gr.Tab('Questions'):
chat = gr.Chatbot(label='Conversation', height=460)
question = gr.Textbox(label='Your question', placeholder='What does this document say about…?')
with gr.Row():
ask_button = gr.Button('Ask', variant='primary')
reset = gr.Button('New conversation')
with gr.Tab('Summaries'):
gr.Markdown('Summaries cover each selected document independently of your conversation.')
summary_button = gr.Button('Summarize documents', variant='primary')
summary = gr.Markdown()
with gr.Accordion('Extracted text', open=False):
text_doc = gr.Dropdown(choices=[], label='Document')
transcript = gr.Textbox(label='Original text', lines=18, interactive=False)
gr.Markdown('Check important findings against the original PDF.')
files.change(read_uploads, files, [status, selected, text_doc, chat, summary, transcript])
text_doc.change(extracted_text, text_doc, transcript)
for trigger in (ask_button.click, question.submit):
trigger(ask, [question, selected, language], [chat, question])
reset.click(clear_chat, outputs=chat)
summary_button.click(summarize, [selected, language], summary)
demo.load(initialize, outputs=lifecycle)
demo.unload(unload)
return demo.queue(default_concurrency_limit=1)
if __name__ == "__main__":
logging.info("LexQuery Gemini model: %s", GEMINI_MODEL)
build_app().launch(
server_name="0.0.0.0",
server_port=7860,
share=False,
) |