Spaces:
Running
Running
| # Full LLM translation and processing combined with BM25 which helps idenitfy key words. | |
| import pytesseract | |
| from pdf2image import convert_from_path | |
| import requests | |
| import json | |
| import time | |
| import asyncio | |
| import pypdf | |
| import re | |
| import hashlib | |
| import base64 | |
| from io import BytesIO | |
| ollama_semaphore = asyncio.Semaphore(1) | |
| cloud_lock = asyncio.Lock() | |
| LAST_CLOUD_CALL_TIME = 0.0 | |
| CLOUD_DELAY_SECONDS = 13.0 | |
| def get_file_hash(filepath): | |
| h = hashlib.sha256() | |
| with open(filepath, 'rb') as file: | |
| while chunk := file.read(8192): | |
| h.update(chunk) | |
| return h.hexdigest() | |
| def extract_text_from_pdf(doc_path): | |
| extracted_chunks = [] | |
| try: | |
| reader = pypdf.PdfReader(doc_path) | |
| for page_idx, page in enumerate(reader.pages): | |
| text = page.extract_text() | |
| if text and len(text.strip()) > 50: | |
| _chunk_text_layout_aware(text, page_idx, extracted_chunks) | |
| else: | |
| images = convert_from_path(doc_path, first_page=page_idx+1, last_page=page_idx+1) | |
| if images: | |
| ocr_text = pytesseract.image_to_string(images[0], lang="eng+spa+ara+chi_sim+ita+msa") | |
| _chunk_text_layout_aware(ocr_text, page_idx, extracted_chunks) | |
| except Exception as e: | |
| print(f"Extraction error: {e}") | |
| return extracted_chunks | |
| def _chunk_text_layout_aware(text, page_idx, extracted_chunks): | |
| text = re.sub(r'\n{3,}', '\n\n', text) | |
| blocks = text.split('\n\n') | |
| current_chunk = "" | |
| for block in blocks: | |
| if len(current_chunk) + len(block) < 800: | |
| current_chunk += block + "\n\n" | |
| else: | |
| if current_chunk.strip(): | |
| extracted_chunks.append({"page": page_idx + 1, "text": current_chunk.strip()}) | |
| current_chunk = current_chunk[-150:] + block + "\n\n" | |
| if current_chunk.strip(): | |
| extracted_chunks.append({"page": page_idx + 1, "text": current_chunk.strip()}) | |
| def get_page_image_b64(doc_path, page_num): | |
| images = convert_from_path(doc_path, first_page=page_num, last_page=page_num) | |
| if not images: return None | |
| buffered = BytesIO() | |
| images[0].save(buffered, format="JPEG", quality=85) | |
| return base64.b64encode(buffered.getvalue()).decode("utf-8") | |
| async def expand_query_async(query, document_anchor, api_key=""): | |
| """Dynamically detects document language via Anchor and expands the query.""" | |
| sys_prompt = ( | |
| "You are an expert multilingual legal translator. " | |
| "1. Read the document excerpt to detect its language and jurisdiction.\n" | |
| "2. Translate the 'Target Field' into 5 exact synonyms used in that specific language context.\n" | |
| "3. Output ONLY the synonyms separated by spaces. No other text." | |
| ) | |
| user_prompt = f"Document Excerpt: {document_anchor[:1000]}\n\nTarget Field: {query}" | |
| try: | |
| if api_key: | |
| url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={api_key}" | |
| payload = {"contents": [{"parts": [{"text": f"{sys_prompt}\n\n{user_prompt}"}]}]} | |
| res = await asyncio.to_thread(requests.post, url, headers={"Content-Type": "application/json"}, json=payload, timeout=10) | |
| return res.json()['candidates'][0]['content']['parts'][0]['text'].strip() | |
| return query | |
| except: | |
| return query | |
| async def query_llm_async(model_name, system_prompt, user_prompt, response_type="extraction", api_key="", doc_path=None, page_nums=None): | |
| """CLOUD-ONLY ROUTER""" | |
| if not api_key: | |
| return "Error: Gemini API Key Required", 0.0 | |
| # Set schemas | |
| if response_type == "extraction": | |
| schema = { | |
| "type": "OBJECT", | |
| "properties": { | |
| "step_1_evidence": {"type": "STRING"}, | |
| "step_2_math_and_logic": {"type": "STRING"}, | |
| "extracted_value": {"type": "STRING"} | |
| }, | |
| "required": ["step_1_evidence", "step_2_math_and_logic", "extracted_value"] | |
| } | |
| target_key = "extracted_value" | |
| else: | |
| schema = { | |
| "type": "OBJECT", | |
| "properties": { | |
| "internal_calculations_do_not_show_user": {"type": "STRING"}, | |
| "final_response": {"type": "STRING"} | |
| }, | |
| "required": ["internal_calculations_do_not_show_user", "final_response"] | |
| } | |
| target_key = "final_response" | |
| # Rate limiting lock | |
| global LAST_CLOUD_CALL_TIME | |
| async with cloud_lock: | |
| current_time = time.time() | |
| if current_time - LAST_CLOUD_CALL_TIME < CLOUD_DELAY_SECONDS: | |
| await asyncio.sleep(CLOUD_DELAY_SECONDS - (current_time - LAST_CLOUD_CALL_TIME)) | |
| LAST_CLOUD_CALL_TIME = time.time() | |
| url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent?key={api_key}" | |
| parts = [{"text": f"SYSTEM INSTRUCTIONS:\n{system_prompt}\n\nUSER PROMPT:\n{user_prompt}"}] | |
| if doc_path and page_nums: | |
| parts[0]["text"] += "\n\nLook at the provided document images to find the exact value." | |
| for p_num in page_nums[:2]: | |
| b64_img = await asyncio.to_thread(get_page_image_b64, doc_path, p_num) | |
| if b64_img: parts.append({"inlineData": {"mimeType": "image/jpeg", "data": b64_img}}) | |
| payload = { | |
| "contents": [{"parts": parts}], | |
| "generationConfig": { | |
| "temperature": 0.0, | |
| "responseMimeType": "application/json", | |
| "responseSchema": schema | |
| } | |
| } | |
| start_time = time.time() | |
| try: | |
| response = await asyncio.to_thread(requests.post, url, headers={"Content-Type": "application/json"}, json=payload, timeout=90) | |
| if response.status_code == 429: | |
| return "RATE_LIMIT_EXCEEDED", round(time.time() - start_time, 2) | |
| response.raise_for_status() | |
| raw_content = response.json()['candidates'][0]['content']['parts'][0]['text'] | |
| latency = round(time.time() - start_time, 2) | |
| try: | |
| return str(json.loads(raw_content).get(target_key, "Not Found")).strip(), latency | |
| except: | |
| return "Parse Error", latency | |
| except Exception as e: | |
| return f"API Error", round(time.time() - start_time, 2) |