用Embeddings+Reranking+LLM构建可靠的内容分类流水线

dev.to · AI#RAG#Embeddings#LLM架构

Editor brief · 编辑速览

from __future__ import annotations import json import os import time import urllib.error import urllib.request from dataclasses import dataclass from typing import Any API_URL = os.environ["INFRAI_BASE_URL"].rstrip("/") + "/v1/chat/completions" ALLOWED_TOPICS = {"spam", "harassment", "privacy", "other"} @dataclass(frozen=True) class Classification: topic: str confidence: float evidence_ids: tuple[str, ...] def post_json(payload: dict[str, Any], attempts: int = 4) -> dict[str, Any]: api_key = os.environ["INFRAI_API_KEY"] body = json.dumps(payload).encode("utf-8") for attempt in range(attempts): request = urllib.request.Request( API_URL, data=body, headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", }, method="POST", ) try: with urllib.request.urlopen(request, timeout=30) as response: return json.loads(response.read()) except urllib.error.HTTPError as error: reason = error.read().decode("utf-8", errors="replace") if error.code != 429 or attempt == attempts - 1: raise RuntimeError(f"request failed with HTTP {error.code}: {reason}") from error retry_after = error.headers.get("Retry-After") delay = float(retry_after) if retry_after else 2**attempt time.sleep(delay) raise RuntimeError("retry budget exhausted") def validate(raw: dict[str, Any], shown_ids: set[str]) -> Classification: if set(raw) != {"topic", "confidence", "evidence_ids"}: raise ValueError("classifier output has missing or unknown fields") if raw["topic"] not in ALLOWED_TOPICS: raise ValueError("classifier returned an unknown topic") if not isinstance(raw["confidence"], (int, float)) or not 0 <= raw["confidence"] <= 1: raise ValueError("confidence must be a number from 0 through 1") evidence = raw["evidence_ids"] if not isinstance(evidence, list) or not evidence or any(item not in shown_ids for item in evidence): raise ValueError("evidence must identify guidance shown to the classifier") return Classification(raw["topic"], float(raw["confidence"]), tuple(evidence)) def classify(report: str, guidance: list[dict[str, str]]) -> Classification: schema = { "name": "moderation_topic", "strict": True, "schema": { "type": "object", "properties": { "topic": {"type": "string", "enum": sorted(ALLOWED_TOPICS)}, "confidence": {"type": "number", "minimum": 0, "maximum": 1}, "evidence_ids": {"type": "array", "items": {"type": "string"}, "minItems": 1}, }, "required": ["topic", "confidence", "evidence_ids"], "additionalProperties": False, }, } result = post_json( { "model": "auto", "messages": [ {"role": "system", "content": "Classify the report using only the supplied guidance."}, {"role": "user", "content": json.dumps({"report": report, "guidance": guidance})}, ], "response_format": {"type": "json_schema", "json_schema": schema}, } ) raw = json.loads(result["choices"][0]["message"]["content"]) return validate(raw, {item["passage_id"] for item in guidance}) if __name__ == "__main__": label = classify( "Repeated unsolicited promotion sent to a maintainer", [ {"passage_id": "tax-spam-2", "text": "Repeated unsolicited promotion maps to spam."}, {"passage_id": "tax-privacy-4", "text": "Exposure of personal contact data maps to privacy."}, ], ) print(label)

完整中文译文