File size: 9,265 Bytes
b69b2d4 | 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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | import torch
from sentence_transformers import SentenceTransformer
from typing import List, Optional
from .model import ICEClassifier
from .di3 import run_di3
from .schemas import ClassificationResult
from sqlalchemy.orm import Session
class PyTorchClassifier:
def __init__(self, model_path="models/classifier/ice_classifier.pt",
schema_path="data/labeled/label_schema.json"):
# These lists are fixed – the order must match training
self.TOPIC_LABELS = [
"Software_&_Tech", "STEM_&_Academics", "Business_&_Finance",
"Creative_&_Media", "Admin_&_Productivity", "Lifestyle_&_Health",
"Social_&_Relationships", "World_&_Current_Events", "Meta_AI",
"Null_Noise", "General_Reference_&_Trivia"
]
self.INTENT_LABELS = [
"Factual_Retrieval", "Troubleshooting", "Generation", "Ideation",
"Analysis_&_Summarization", "Strategic_Planning", "Decision_Making",
"Emotional_Processing", "Utility_Formatting", "Casual_Banter",
"Open_Exploration"
]
self.CONTEXT_RELIANCE_LABELS = [
"Zero_Shot", "Long_Term_Memory", "Real_Time_Search"
]
# Load model on CPU
self.model = ICEClassifier()
self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
self.model.eval()
# Embedder also on CPU
# Embedder also on CPU – Qwen3-Embedding truncated to 384 dim for compatibility
self.embedder = SentenceTransformer(
"Qwen/Qwen3-Embedding-0.6B",
device="cpu",
truncate_dim=384
)
def _get_context_turns(self, conversation_id: str, n: int = 3, max_total_words: int = 500) -> str:
"""Return a truncated, summary‑preferring context string from the last *n* turns."""
# Local import to avoid circular dependency at module level
from src.api.db import SessionLocal
db = SessionLocal()
try:
from src.memory.models import EpisodicMemory
turns = (
db.query(EpisodicMemory)
.filter_by(conversation_id=conversation_id)
.order_by(EpisodicMemory.timestamp.desc())
.limit(n)
.all()
)
turns.reverse()
parts = []
total_words = 0
for t in turns:
# Prefer summary, fall back to raw text (truncated)
text = t.summary_text or ""
if not text and t.raw_text:
words = t.raw_text.split()
text = " ".join(words[:150]) + "…" if len(words) > 150 else t.raw_text
if not text:
continue
word_count = len(text.split())
if total_words + word_count > max_total_words:
remaining = max_total_words - total_words
if remaining > 20:
w = text.split()
text = " ".join(w[:remaining]) + "…"
parts.append(text)
break
parts.append(text)
total_words += word_count
return "\n".join(parts)
finally:
db.close()
# ------------------------------------------------------------------
# Main entry point
# ------------------------------------------------------------------
def classify(
self,
prompt: str,
conversation_history: Optional[List[str]] = None,
conversation_length: int = 0,
conversation_id: Optional[str] = None,
) -> ClassificationResult:
"""Public entry point. Runs DI3 first, falls back to ML.
When *conversation_id* is given, the last 3 turns are used as context
(auto‑truncated) to improve the ML classifier's accuracy.
"""
if conversation_history is None:
conversation_history = []
di3_result = run_di3(prompt, conversation_length, conversation_history)
if di3_result is not None:
# If DI3 forced LTM but left topic/intent blank, let the ML
# classifier provide the actual tags while keeping the LTM decision.
if not di3_result.topic_tags or not di3_result.intent_tags:
ml_result = self._run_ml_classifier(prompt, conversation_id)
if di3_result.context_reliance == "Long_Term_Memory":
ml_result.context_reliance = "Long_Term_Memory"
return self._apply_hard_overrides(ml_result, prompt)
else:
return self._apply_hard_overrides(di3_result, prompt)
return self._run_ml_classifier(prompt, conversation_id)
def _run_ml_classifier(self, prompt: str, conversation_id: Optional[str] = None) -> ClassificationResult:
"""Original ML classification path (now private)."""
with torch.no_grad():
# Build context text if conversation_id is available
context_text = None
if conversation_id:
try:
context_text = self._get_context_turns(conversation_id)
except Exception:
context_text = None
if context_text:
prefixed_prompt = (
f"Conversation context (summarized):\n{context_text}\n\n"
f"Given the above conversation and the user's latest prompt, "
f"predict:\n"
f"1. TOPIC: what is the subject (Software_&_Tech, Creative_&_Media, etc.)\n"
f"2. INTENT: what is the user trying to do (Factual_Retrieval, Troubleshooting, etc.)\n"
f"3. CONTEXT RELIANCE: does the user need memory (Zero_Shot, Long_Term_Memory, Real_Time_Search)\n\n"
f"User prompt: {prompt}"
)
else:
prefixed_prompt = (
f"Given a user prompt, predict:\n"
f"1. TOPIC: what is the subject (Software_&_Tech, Creative_&_Media, etc.)\n"
f"2. INTENT: what is the user trying to do (Factual_Retrieval, Troubleshooting, etc.)\n"
f"3. CONTEXT RELIANCE: does the user need memory (Zero_Shot, Long_Term_Memory, Real_Time_Search)\n\n"
f"User prompt: {prompt}"
)
embedding = self.embedder.encode(prefixed_prompt, convert_to_tensor=True).unsqueeze(0).float()
outputs = self.model(embedding) # (1, 25)
topic_out = outputs[:, :11] # (1, 11)
intent_out = outputs[:, 11:22] # (1, 11)
ctx_out = outputs[:, 22:] # (1, 3)
topic_probs = torch.sigmoid(topic_out).squeeze(0) # (11,)
intent_probs = torch.sigmoid(intent_out).squeeze(0) # (11,)
ctx_probs = torch.softmax(ctx_out, dim=1).squeeze(0) # (3,)
# Build tag lists
topic_tags = [self.TOPIC_LABELS[i] for i in range(len(self.TOPIC_LABELS))
if topic_probs[i] > 0.3]
intent_tags = [self.INTENT_LABELS[i] for i in range(len(self.INTENT_LABELS))
if intent_probs[i] > 0.3]
if not topic_tags:
topic_tags = [self.TOPIC_LABELS[torch.argmax(topic_probs).item()]]
if not intent_tags:
intent_tags = [self.INTENT_LABELS[torch.argmax(intent_probs).item()]]
context_reliance = self.CONTEXT_RELIANCE_LABELS[torch.argmax(ctx_probs).item()]
# Combine probabilities
raw_probs = topic_probs.tolist() + intent_probs.tolist() + ctx_probs.tolist()
max_confidence = max(raw_probs)
result = ClassificationResult(
topic_tags=topic_tags,
intent_tags=intent_tags,
context_reliance=context_reliance,
raw_probs=raw_probs,
max_confidence=max_confidence,
prompt=prompt,
)
return self._apply_hard_overrides(result, prompt)
def _apply_hard_overrides(
self, result: ClassificationResult, prompt: str
) -> ClassificationResult:
"""Apply creative/software LTM overrides, but never downgrade an
existing Long_Term_Memory decision (e.g. from DI3 or LTM bias)."""
# If LTM has already been enforced (by DI3 or API‑level bias), keep it
if result.context_reliance == "Long_Term_Memory":
return result
if "Creative_&_Media" in result.topic_tags:
result.context_reliance = "Long_Term_Memory"
if "Software_&_Tech" in result.topic_tags:
referential_words = [
"my", "our", "mine", "ours", "we", "us",
"this", "that", "these", "those", "the",
"it", "they", "them", "their",
"previous", "last", "before", "yesterday", "earlier",
"again", "still", "same",
]
prompt_lower = prompt.lower()
if any(word in prompt_lower for word in referential_words):
result.context_reliance = "Long_Term_Memory"
return result |