File size: 2,226 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 | """ICE v2 learned-head input and decoding, without database or DI3 rules."""
import json
from pathlib import Path
import torch
def build_input(prompt: str, context_text: str | None = None) -> str:
if context_text:
return (
f"Conversation context (summarized):\n{context_text}\n\n"
"Given the above conversation and the user's latest prompt, predict:\n"
"1. TOPIC: what is the subject (Software_&_Tech, Creative_&_Media, etc.)\n"
"2. INTENT: what is the user trying to do (Factual_Retrieval, Troubleshooting, etc.)\n"
"3. CONTEXT RELIANCE: does the user need memory (Zero_Shot, Long_Term_Memory, Real_Time_Search)\n\n"
f"User prompt: {prompt}"
)
return (
"Given a user prompt, predict:\n"
"1. TOPIC: what is the subject (Software_&_Tech, Creative_&_Media, etc.)\n"
"2. INTENT: what is the user trying to do (Factual_Retrieval, Troubleshooting, etc.)\n"
"3. CONTEXT RELIANCE: does the user need memory (Zero_Shot, Long_Term_Memory, Real_Time_Search)\n\n"
f"User prompt: {prompt}"
)
def predict_head(model, embedder, prompt: str, context_text: str | None = None):
config = json.loads(Path(__file__).with_name("config.json").read_text())
embedding = embedder.encode(build_input(prompt, context_text), convert_to_tensor=True)
if embedding.shape != (384,):
raise ValueError("ICE v2 requires exactly 384 embedding coordinates")
with torch.no_grad():
outputs = model(embedding.unsqueeze(0).float())
topic = torch.sigmoid(outputs[0, :11])
intent = torch.sigmoid(outputs[0, 11:22])
context = torch.softmax(outputs[0, 22:], dim=0)
def tags(probs, labels):
return [labels[i] for i, p in enumerate(probs) if p > 0.3] or [labels[probs.argmax().item()]]
probabilities = topic.tolist() + intent.tolist() + context.tolist()
return {
"topic_tags": tags(topic, config["topic_labels"]),
"intent_tags": tags(intent, config["intent_labels"]),
"context_reliance": config["context_labels"][context.argmax().item()],
"raw_probs": probabilities,
"max_confidence": max(probabilities),
}
|