diff --git a/src/llm/chunked_extraction.py b/src/llm/chunked_extraction.py index 6c15275..e1d9e3f 100644 --- a/src/llm/chunked_extraction.py +++ b/src/llm/chunked_extraction.py @@ -3,7 +3,6 @@ import sys import gc import json -import joblib import xgboost as xgb from pathlib import Path from collections import Counter @@ -13,21 +12,7 @@ sys.path.insert(0, str(project_root)) from src.llm.llm_client import extract_metrics_from_text - - -def load_classifier(model_dir="src/model/models"): - """Load XGBoost classifier and TF-IDF vectorizer.""" - model_path = Path(model_dir) / "pdf_classifier.json" - vectorizer_path = Path(model_dir) / "tfidf_vectorizer.pkl" - - if not model_path.exists(): - raise FileNotFoundError(f"Model not found: {model_path}") - - model = xgb.Booster() - model.load_model(str(model_path)) - vectorizer = joblib.load(vectorizer_path) - - return model, vectorizer +from src.model.pdf_classifier import load_classifier def chunk_text(text, chunk_size=3000, overlap=300): @@ -100,7 +85,7 @@ def extract_with_chunking( text, model_dir="src/model/models", llm_model="qwen2.5:7b", # Changed from biomistral - num_ctx=2048, + num_ctx=8192, top_n=3, chunk_size=3000, overlap=300, @@ -108,7 +93,7 @@ def extract_with_chunking( """Main extraction with chunking pipeline.""" print(f" [CHUNK] Loading classifier...", file=sys.stderr) - model, vectorizer = load_classifier(model_dir) + model, vectorizer, _encoder = load_classifier(model_dir) chunks = chunk_text(text, chunk_size, overlap) print(f" [CHUNK] Split into {len(chunks)} chunks", file=sys.stderr) @@ -139,7 +124,7 @@ def extract_with_chunking( try: metrics = extract_metrics_from_text( - text=chunk[:2500], + text=chunk, model=llm_model, num_ctx=num_ctx, ) diff --git a/src/llm/llm_client.py b/src/llm/llm_client.py index cd25900..1103c97 100644 --- a/src/llm/llm_client.py +++ b/src/llm/llm_client.py @@ -16,6 +16,8 @@ import logging import re import sys +import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError from pathlib import Path from ollama import chat @@ -29,6 +31,32 @@ log = logging.getLogger(__name__) +_OLLAMA_TIMEOUT = 120 # seconds +_MAX_RETRIES = 3 +_BACKOFF_BASE = 2 # seconds + + +def _call_ollama_with_retry(model, messages, format, options): + """Call Ollama with timeout and exponential backoff on transient failures.""" + transient_errors = (ConnectionError, OSError, TimeoutError, FuturesTimeoutError) + + for attempt in range(_MAX_RETRIES): + try: + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(chat, messages=messages, model=model, format=format, options=options) + return future.result(timeout=_OLLAMA_TIMEOUT) + except FuturesTimeoutError: + log.warning("Ollama call timed out (attempt %d/%d)", attempt + 1, _MAX_RETRIES) + except transient_errors as e: + log.warning("Transient Ollama error (attempt %d/%d): %s", attempt + 1, _MAX_RETRIES, e) + + if attempt < _MAX_RETRIES - 1: + wait = _BACKOFF_BASE**attempt + log.info("Retrying in %ds...", wait) + time.sleep(wait) + + raise RuntimeError(f"Ollama call failed after {_MAX_RETRIES} attempts") + def extract_metrics_from_text( text: str, @@ -125,9 +153,9 @@ def extract_metrics_from_text( TEXT {text} """ - response = chat( - messages=[{"role": "user", "content": prompt}], + response = _call_ollama_with_retry( model=model, + messages=[{"role": "user", "content": prompt}], format=PredatorDietMetrics.model_json_schema(), options={"num_ctx": num_ctx}, ) @@ -184,9 +212,9 @@ def extract_metrics_from_text( retry_prompt += _hints.get(field, "") retry_prompt += f"\nTEXT\n{text}" - retry_response = chat( - messages=[{"role": "user", "content": retry_prompt}], + retry_response = _call_ollama_with_retry( model=model, + messages=[{"role": "user", "content": retry_prompt}], format=PredatorDietMetrics.model_json_schema(), options={"num_ctx": num_ctx}, )