Spaces:
Build error
Build error
File size: 5,212 Bytes
8c3e275 | 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 | from __future__ import annotations
from pageparse.config import settings
from pageparse.schema import SourceParseResult, Record
SCHEMA_DETECTION_KEYWORDS = {
"todo": ["todo", "to-do", "task", "deadline", "due", "priority", "action item"],
"meeting": ["meeting", "agenda", "speaker", "minutes", "attendee", "discussion"],
"recipe": ["recipe", "ingredient", "tbsp", "tsp", "cup", "cook", "bake", "preheat"],
"survey": ["survey", "questionnaire", "question", "response", "feedback"],
"notes": ["note", "lecture", "chapter", "section", "topic"],
}
DEFAULT_SCHEMA_PROMPTS = {
"meeting": (
"Extract records of type 'action_item', 'note', or 'transcript_segment' from the following text. "
"Include speaker, timestamp, priority, and due_date where present. "
"Do NOT change any words. Preserve original text exactly."
),
"recipe": (
"Extract ingredients as type 'key_value' (category 'Ingredient') and "
"steps as type 'note' (category 'Instruction'). "
"Do NOT change any words. Preserve original text exactly."
),
"survey": (
"Extract questions as type 'note' (category 'Question') and "
"responses as type 'key_value' (category 'Response'). "
"Do NOT change any words. Preserve original text exactly."
),
"notes": (
"Extract key topics as type 'note' with category from the following notes. "
"Include any action items with priority and due_date where present. "
"Do NOT change any words. Preserve original text exactly."
),
}
class Extractor:
def __init__(self) -> None:
self._grammar: str | None = None
grammar_path = settings.grammar_path()
if grammar_path.exists():
self._grammar = grammar_path.read_text(encoding="utf-8")
def detect_schema_type(self, raw_text: str) -> str:
if not settings.auto_schema:
return "todo"
text_lower = raw_text.lower()
scores: dict[str, int] = {}
for schema_type, keywords in SCHEMA_DETECTION_KEYWORDS.items():
score = sum(text_lower.count(kw) for kw in keywords)
if score > 0:
scores[schema_type] = score
if not scores:
return "todo"
return max(scores, key=scores.get)
def extract(
self, raw_text: str, source_file: str, schema_type: str | None = None
) -> SourceParseResult:
if schema_type is None or schema_type == "auto":
schema_type = self.detect_schema_type(raw_text)
result = SourceParseResult(
source_file=source_file,
source_type="image",
records=[],
)
if not raw_text.strip() or raw_text in ("[UNCLEAR]", "[inaudible]"):
return result
result.records = self._line_based_fallback(raw_text)
if self._grammar and raw_text.strip():
try:
structured = self._ollama_structured_extract(raw_text, source_file, schema_type)
if structured and structured.records:
result = structured
result.source_file = source_file
except Exception:
pass
return result
def _ollama_structured_extract(
self, raw_text: str, source_file: str, schema_type: str
) -> SourceParseResult | None:
import json
import urllib.request
user_msg = DEFAULT_SCHEMA_PROMPTS.get(
schema_type,
"Extract tasks with due_date, priority, status, and category from the following text. "
"Do NOT change any words. Preserve original text exactly."
)
system_msg = "You are a structured data extractor. Output only valid JSON matching the schema."
prompt = f"{system_msg}\n\n{user_msg}\n\nExtracted Text:\n{raw_text}\n\nSource File: {source_file}"
payload = {
"model": settings.slm_model,
"prompt": prompt,
"stream": False,
"options": {"temperature": 0.1},
}
url = f"{settings.ollama_url}/api/generate"
data = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
with urllib.request.urlopen(req, timeout=30) as response:
body = response.read().decode("utf-8")
res = json.loads(body)
raw = res.get("response", "").strip()
if raw:
try:
return SourceParseResult.model_validate_json(raw)
except Exception:
pass
return None
def _line_based_fallback(self, raw_text: str) -> list[Record]:
records: list[Record] = []
lines = raw_text.strip().split("\n")
for line in lines:
line = line.strip()
if not line:
continue
records.append(
Record(
type="note",
content=line,
confidence=0.95,
)
)
if not records:
records.append(Record(type="note", content=raw_text.strip()))
return records
|