128 lines
5.2 KiB
Python
128 lines
5.2 KiB
Python
"""SHARED CONTRACT — the prompt + output schema the model is trained and served on.
|
|
|
|
This file is the single source of truth for:
|
|
- FIELD_ORDER : canonical key order of the extracted JSON
|
|
- build_prompt() : the EXACT prompt string used at train AND inference time
|
|
- format_completion() : training target serialization (also used by eval)
|
|
- extract_json() : brace-matched JSON extraction from raw model text
|
|
- HotelExtraction : Pydantic schema for validation
|
|
- validate() : parse+validate -> clean dict (or None)
|
|
|
|
It is intentionally DEPENDENCY-LIGHT (only json + pydantic) so both the heavy
|
|
training env and the lean serving env can import it. A COPY of this file lives in
|
|
BOTH training/ and deployment/; tests/test_contract_parity.py asserts they are
|
|
byte-identical, so the prompt can never silently drift between train and serve.
|
|
|
|
DO NOT EDIT one copy without the other. Edit the source, re-sync, re-run the test.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any, Literal, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
|
|
|
DEFAULT_BASE_MODEL = "Qwen/Qwen3-1.7B"
|
|
|
|
# `filters` is emitted as a comma-separated string of NATURAL PHRASES
|
|
# (e.g. "swimming pool, pet friendly"), NOT production codes. A downstream matcher
|
|
# resolves phrases -> codes (FL_HF_29, ...), keeping the model taxonomy-agnostic.
|
|
FIELD_ORDER = [
|
|
"destination", "locality", "hotelName", "checkinDate", "checkoutDate",
|
|
"adultCount", "roomCount", "childCount", "childAges", "infantCount",
|
|
"sortCriteria", "minPrice", "maxPrice", "filters", "deepSearch", "isNearMe", "resetAction",
|
|
]
|
|
|
|
|
|
def build_prompt(query: str, today: str) -> str:
|
|
"""The exact user-message content used for BOTH training and inference."""
|
|
return (
|
|
"Extract hotel-search entities.\n"
|
|
"Return strict JSON only.\n"
|
|
"Schema:\n"
|
|
"destination, locality, hotelName, checkinDate, checkoutDate, "
|
|
"adultCount, roomCount, childCount, childAges, infantCount, "
|
|
"sortCriteria, minPrice, maxPrice, filters, deepSearch, isNearMe, resetAction.\n\n"
|
|
"Rules:\n"
|
|
"- Dates are DDMMYYYY.\n"
|
|
'- deepSearch and isNearMe must be "true" or "false".\n'
|
|
"- filters is a comma-separated list of amenity/type phrases (e.g. "
|
|
'"swimming pool, pet friendly"); prefix removals with "no ".\n'
|
|
"- Omit unknown optional fields.\n\n"
|
|
f"today={today}\n"
|
|
f"query={query}"
|
|
)
|
|
|
|
|
|
def format_completion(expected: dict[str, Any]) -> str:
|
|
"""Stable-key-order minified JSON — the training target / gold serialization."""
|
|
ordered = {f: expected[f] for f in FIELD_ORDER if f in expected}
|
|
for k in expected:
|
|
if k not in ordered:
|
|
ordered[k] = expected[k]
|
|
return json.dumps(ordered, ensure_ascii=False, separators=(",", ":"))
|
|
|
|
|
|
def extract_json(text: str) -> dict | None:
|
|
"""Pull the first balanced {...} object out of raw model text."""
|
|
start = text.find("{")
|
|
if start < 0:
|
|
return None
|
|
depth = 0
|
|
in_str = esc = False
|
|
for i in range(start, len(text)):
|
|
c = text[i]
|
|
if esc:
|
|
esc = False
|
|
continue
|
|
if c == "\\":
|
|
esc = True
|
|
continue
|
|
if c == '"':
|
|
in_str = not in_str
|
|
continue
|
|
if in_str:
|
|
continue
|
|
if c == "{":
|
|
depth += 1
|
|
elif c == "}":
|
|
depth -= 1
|
|
if depth == 0:
|
|
try:
|
|
return json.loads(text[start:i + 1])
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return None
|
|
|
|
|
|
class HotelExtraction(BaseModel):
|
|
"""Output contract. `filters` is a comma-separated PHRASE string (e.g.
|
|
"swimming pool, pet friendly"); a downstream matcher resolves it to codes."""
|
|
model_config = ConfigDict(extra="forbid")
|
|
deepSearch: Literal["true", "false"]
|
|
isNearMe: Literal["true", "false"]
|
|
destination: Optional[str] = Field(default=None, min_length=1)
|
|
locality: Optional[str] = Field(default=None, min_length=1)
|
|
hotelName: Optional[str] = Field(default=None, min_length=1)
|
|
checkinDate: Optional[str] = Field(default=None, pattern=r"^\d{8}$")
|
|
checkoutDate: Optional[str] = Field(default=None, pattern=r"^\d{8}$")
|
|
adultCount: Optional[int] = Field(default=None, ge=0, le=20)
|
|
roomCount: Optional[int] = Field(default=None, ge=0, le=20)
|
|
childCount: Optional[int] = Field(default=None, ge=0, le=20)
|
|
childAges: Optional[list[int]] = None
|
|
infantCount: Optional[int] = Field(default=None, ge=0, le=20)
|
|
sortCriteria: Optional[Literal["SC_P_LH", "SC_P_HL", "SC_UR", "SC_P", "SC_DIST"]] = None
|
|
minPrice: Optional[int] = Field(default=None, ge=0, le=10_000_000)
|
|
maxPrice: Optional[int] = Field(default=None, ge=0, le=10_000_000)
|
|
filters: Optional[str] = Field(default=None, min_length=1)
|
|
resetAction: Optional[Literal["filters", "guests", "dates", "all"]] = None
|
|
|
|
|
|
def validate(d: dict) -> dict | None:
|
|
"""Validate a parsed dict against the schema; return clean dict or None."""
|
|
try:
|
|
return HotelExtraction.model_validate(d).model_dump(exclude_none=True)
|
|
except ValidationError:
|
|
return None
|