summaryrefslogtreecommitdiff
path: root/src/personalization/serving
diff options
context:
space:
mode:
Diffstat (limited to 'src/personalization/serving')
-rw-r--r--src/personalization/serving/__init__.py22
-rw-r--r--src/personalization/serving/api/__init__.py0
-rw-r--r--src/personalization/serving/api/main.py0
-rw-r--r--src/personalization/serving/api/routes/__init__.py0
-rw-r--r--src/personalization/serving/api/routes/feedback.py0
-rw-r--r--src/personalization/serving/api/routes/query.py0
-rw-r--r--src/personalization/serving/api/routes/users.py0
-rw-r--r--src/personalization/serving/api/schemas.py0
-rw-r--r--src/personalization/serving/personalized_llm.py837
9 files changed, 859 insertions, 0 deletions
diff --git a/src/personalization/serving/__init__.py b/src/personalization/serving/__init__.py
new file mode 100644
index 0000000..11adcf8
--- /dev/null
+++ b/src/personalization/serving/__init__.py
@@ -0,0 +1,22 @@
+# Personalization Serving Module
+#
+# This module provides the interface layer for the personalization system.
+
+from personalization.serving.personalized_llm import (
+ PersonalizedLLM,
+ AssistantResponse,
+ UsageStats,
+ DebugInfo,
+ Feedback,
+ create_personalized_llm,
+)
+
+__all__ = [
+ "PersonalizedLLM",
+ "AssistantResponse",
+ "UsageStats",
+ "DebugInfo",
+ "Feedback",
+ "create_personalized_llm",
+]
+
diff --git a/src/personalization/serving/api/__init__.py b/src/personalization/serving/api/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/__init__.py
diff --git a/src/personalization/serving/api/main.py b/src/personalization/serving/api/main.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/main.py
diff --git a/src/personalization/serving/api/routes/__init__.py b/src/personalization/serving/api/routes/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/routes/__init__.py
diff --git a/src/personalization/serving/api/routes/feedback.py b/src/personalization/serving/api/routes/feedback.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/routes/feedback.py
diff --git a/src/personalization/serving/api/routes/query.py b/src/personalization/serving/api/routes/query.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/routes/query.py
diff --git a/src/personalization/serving/api/routes/users.py b/src/personalization/serving/api/routes/users.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/routes/users.py
diff --git a/src/personalization/serving/api/schemas.py b/src/personalization/serving/api/schemas.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/src/personalization/serving/api/schemas.py
diff --git a/src/personalization/serving/personalized_llm.py b/src/personalization/serving/personalized_llm.py
new file mode 100644
index 0000000..2c4d5a8
--- /dev/null
+++ b/src/personalization/serving/personalized_llm.py
@@ -0,0 +1,837 @@
+#!/usr/bin/env python3
+"""
+Personalized LLM Interface for Evaluation.
+
+This module provides the `PersonalizedLLM` class that wraps the entire
+personalization system into a clean interface for evaluation frameworks
+and user simulators.
+
+Interface contract:
+- chat(user_id, query) -> AssistantResponse: Main online interface
+- reset_session(user_id): Clear session history and short-term state
+- reset_user(user_id): Completely reset user (long-term, short-term, memories)
+- apply_feedback(feedback): Apply external feedback for RL updates
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import yaml
+
+# Ensure src is in path for standalone usage
+_src_path = os.path.join(os.path.dirname(__file__), "../../..")
+if _src_path not in sys.path:
+ sys.path.insert(0, _src_path)
+
+from personalization.config.settings import load_local_models_config
+from personalization.config.registry import get_preference_extractor, get_chat_model
+from personalization.models.embedding.qwen3_8b import Qwen3Embedding8B
+from personalization.models.reranker.qwen3_reranker import Qwen3Reranker
+from personalization.user_model.tensor_store import UserTensorStore, UserState
+from personalization.user_model.session_state import OnlineSessionState
+from personalization.user_model.features import ItemProjection
+from personalization.retrieval.preference_store.schemas import (
+ MemoryCard, ChatTurn, PreferenceList, Preference
+)
+from personalization.retrieval.pipeline import retrieve_with_policy, retrieve_no_policy
+from personalization.feedback.handlers import eval_step
+from personalization.user_model.policy.reinforce import reinforce_update_user_state
+
+
+# =============================================================================
+# Data Classes for Interface
+# =============================================================================
+
+@dataclass
+class UsageStats:
+ """Token usage statistics from a chat completion."""
+ prompt_tokens: int
+ completion_tokens: int
+ total_tokens: int
+ model: str
+
+
+@dataclass
+class DebugInfo:
+ """
+ Debug information for analysis and ablation studies.
+ All fields are optional - fill what you have, leave empty what you don't.
+ """
+ selected_memory_ids: List[str] = field(default_factory=list)
+ selected_memory_notes: List[str] = field(default_factory=list)
+ selected_memory_scores: List[float] = field(default_factory=list)
+ user_vector_before: Optional[List[float]] = None
+ user_vector_after: Optional[List[float]] = None
+ extracted_preferences: List[Dict[str, Any]] = field(default_factory=list)
+ extra: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class AssistantResponse:
+ """Response from the personalized LLM chat interface."""
+ answer: str
+ usage: UsageStats
+ debug: Optional[DebugInfo] = None
+
+
+@dataclass
+class Feedback:
+ """
+ Feedback data structure for RL updates from user simulator or judge.
+
+ Attributes:
+ user_id: The user this feedback is for.
+ turn_id: The turn this feedback refers to (from the previous turn).
+ reward: Reward scalar computed by user simulator / judge.
+ gating: Gating flag (1=valid learning signal, 0=skip update).
+ meta: Additional metadata for training/analysis.
+ """
+ user_id: str
+ turn_id: int
+ reward: float
+ gating: float # Can be 0.0 or 1.0, or continuous
+ meta: Dict[str, Any] = field(default_factory=dict)
+
+
+# =============================================================================
+# Internal Session State Extended
+# =============================================================================
+
+@dataclass
+class _SessionContext:
+ """Extended session context for evaluation tracking."""
+ session_state: OnlineSessionState
+ turn_counter: int = 0
+ # Store info needed for apply_feedback
+ pending_rl_update: Optional[Dict[str, Any]] = None
+
+
+# =============================================================================
+# PersonalizedLLM Class
+# =============================================================================
+
+class PersonalizedLLM:
+ """
+ Personalized LLM wrapper for evaluation frameworks.
+
+ This class provides a clean interface that accepts only (user_id, query)
+ for the main chat function, while internally managing:
+ - User state vectors (z_long, z_short)
+ - Session history
+ - Memory retrieval and policy
+ - Preference extraction and storage
+ - RL updates
+
+ Example usage:
+ llm = PersonalizedLLM()
+
+ # Reset user for fresh experiment
+ llm.reset_user("user_123")
+
+ # Start a session
+ llm.reset_session("user_123")
+
+ # Chat
+ response = llm.chat("user_123", "What's a good recipe for dinner?")
+ print(response.answer)
+
+ # Apply feedback from previous turn (from turn 2 onwards)
+ llm.apply_feedback(Feedback(
+ user_id="user_123",
+ turn_id=0,
+ reward=0.8,
+ gating=1.0
+ ))
+ """
+
+ def __init__(
+ self,
+ config_path: Optional[str] = None,
+ user_store_path: str = "data/users/user_store_eval.npz",
+ memory_cards_path: str = "data/corpora/memory_cards.jsonl",
+ memory_embeddings_path: str = "data/corpora/memory_embeddings.npy",
+ item_projection_path: str = "data/corpora/item_projection.npz",
+ only_own_memories: bool = True,
+ enable_preference_extraction: bool = True,
+ enable_rl_updates: bool = True,
+ mode: str = "full", # "full", "nopersonal", or "vanilla"
+ eval_mode: bool = True, # True = greedy selection, False = stochastic sampling
+ device_assignment: Optional[Dict[str, str]] = None, # Multi-GPU support
+ ):
+ """
+ Initialize the PersonalizedLLM.
+
+ Args:
+ config_path: Path to config file. If None, uses default locations.
+ user_store_path: Path to persist user state vectors.
+ memory_cards_path: Path to memory cards JSONL file.
+ memory_embeddings_path: Path to memory embeddings numpy file.
+ item_projection_path: Path to item projection (PCA) file.
+ only_own_memories: If True, only retrieve user's own memories (strict privacy).
+ enable_preference_extraction: If True, extract preferences from user turns.
+ enable_rl_updates: If True, apply RL updates via apply_feedback.
+ mode: "full" for full personalization, "nopersonal" for baseline (no user vector influence),
+ "vanilla" for pure LLM without any memory retrieval or preference extraction.
+ eval_mode: If True, use greedy/deterministic selection (for evaluation).
+ If False, use stochastic sampling (for training/exploration).
+ device_assignment: Optional dict to assign models to specific GPUs.
+ Example: {"embed": "cuda:0", "reranker": "cuda:1", "chat": "cuda:2", "extractor": "cuda:3"}
+ If None, uses "auto" for all models.
+ """
+ self.only_own_memories = only_own_memories
+ self.enable_preference_extraction = enable_preference_extraction
+ self.enable_rl_updates = enable_rl_updates
+ self.mode = mode # "full" or "nopersonal"
+ self.eval_mode = eval_mode # True = greedy, False = sample
+
+ # Multi-GPU device assignment
+ self._device_assignment = device_assignment or {
+ "embed": "auto",
+ "reranker": "auto",
+ "chat": "auto",
+ "extractor": "auto",
+ }
+
+ # Paths
+ self._memory_cards_path = memory_cards_path
+ self._memory_embeddings_path = memory_embeddings_path
+ self._item_projection_path = item_projection_path
+
+ # RL Configuration
+ # Note: beta/eta increased for more significant z_u updates
+ self._rl_cfg = {
+ "item_dim": 256,
+ "beta_long": 2.0, # Increased from 0.1 for stronger personalization
+ "beta_short": 5.0, # Increased from 0.3
+ "tau": 1.0,
+ "eta_long": 0.01, # Increased from 1e-3 for faster learning
+ "eta_short": 0.05, # Increased from 5e-3
+ "ema_alpha": 0.05,
+ "short_decay": 0.1,
+ "dense_topk": 64,
+ "rerank_topk": 3,
+ "max_new_tokens": 512,
+ }
+
+ # Load config and override RL params if available
+ self._load_config(config_path)
+
+ # Load models
+ print("[PersonalizedLLM] Loading models...")
+ self._load_models()
+
+ # Load memory store
+ print("[PersonalizedLLM] Loading memory store...")
+ self._load_memory_store()
+
+ # Initialize user store
+ self._user_store = UserTensorStore(
+ k=self._rl_cfg["item_dim"],
+ path=user_store_path,
+ )
+
+ # Session contexts per user (in-memory)
+ self._sessions: Dict[str, _SessionContext] = {}
+
+ print("[PersonalizedLLM] Initialization complete.")
+
+ def _load_config(self, config_path: Optional[str]):
+ """Load configuration from yaml files."""
+ self._cfg = load_local_models_config()
+
+ # Try to load user_model.yaml for RL params
+ if config_path is None:
+ config_path = "configs/user_model.yaml"
+
+ self._llm_name = "qwen_1_5b" # Default
+
+ try:
+ if os.path.exists(config_path):
+ with open(config_path, "r") as f:
+ user_cfg = yaml.safe_load(f)
+ if user_cfg:
+ # Override RL params if present
+ for key in self._rl_cfg:
+ if key in user_cfg:
+ self._rl_cfg[key] = user_cfg[key]
+ # LLM name
+ if "llm_name" in user_cfg:
+ self._llm_name = user_cfg["llm_name"]
+ except Exception as e:
+ print(f"[PersonalizedLLM] Warning: Failed to load config: {e}")
+
+ def _load_models(self):
+ """Load all ML models with optional multi-GPU assignment."""
+ import torch
+
+ # Report GPU availability
+ num_gpus = torch.cuda.device_count()
+ print(f"[PersonalizedLLM] Available GPUs: {num_gpus}")
+ for i in range(num_gpus):
+ mem = torch.cuda.get_device_properties(i).total_memory / 1e9
+ print(f" GPU {i}: {torch.cuda.get_device_name(i)} ({mem:.1f}GB)")
+
+ embed_device = self._device_assignment.get("embed", "auto")
+ reranker_device = self._device_assignment.get("reranker", "auto")
+ chat_device = self._device_assignment.get("chat", "auto")
+ extractor_device = self._device_assignment.get("extractor", "auto")
+
+ # Embedding model
+ print(f"[PersonalizedLLM] Loading Embedding model on {embed_device}...")
+ self._embed_model = Qwen3Embedding8B(
+ model_path=self._cfg.embedding.qwen3.local_path,
+ dtype=torch.bfloat16,
+ device_map=embed_device,
+ )
+
+ # Reranker
+ print(f"[PersonalizedLLM] Loading Reranker on {reranker_device}...")
+ self._reranker = Qwen3Reranker(
+ model_path=self._cfg.reranker.qwen3_8b.local_path,
+ device_map=reranker_device,
+ dtype=torch.bfloat16,
+ )
+
+ # Chat model (via registry for backend switching)
+ print(f"[PersonalizedLLM] Loading ChatModel: {self._llm_name} on {chat_device}...")
+ # Pass device override if specified (not "auto")
+ device_for_chat = chat_device if chat_device != "auto" else None
+ self._chat_model = get_chat_model(self._llm_name, device_override=device_for_chat)
+
+ # Preference extractor
+ if self.enable_preference_extraction:
+ extractor_name = "qwen3_0_6b_sft"
+ print(f"[PersonalizedLLM] Loading extractor: {extractor_name} on {extractor_device}...")
+ try:
+ self._extractor = get_preference_extractor(extractor_name)
+ except Exception as e:
+ print(f"[PersonalizedLLM] Warning: Failed to load {extractor_name}: {e}. Using rule-based.")
+ self._extractor = get_preference_extractor("rule")
+ else:
+ print("[PersonalizedLLM] Preference extraction disabled, using rule-based extractor.")
+ self._extractor = get_preference_extractor("rule")
+
+ def _load_memory_store(self):
+ """Load memory cards and embeddings."""
+ if not os.path.exists(self._memory_cards_path):
+ print(f"[PersonalizedLLM] Warning: Memory cards not found at {self._memory_cards_path}")
+ self._memory_cards: List[MemoryCard] = []
+ self._memory_embeddings = np.zeros((0, 4096), dtype=np.float32)
+ self._item_vectors = np.zeros((0, self._rl_cfg["item_dim"]), dtype=np.float32)
+ self._projection = None
+ return
+
+ # Load cards
+ self._memory_cards = []
+ with open(self._memory_cards_path, "r") as f:
+ for line in f:
+ line = line.strip()
+ if line:
+ self._memory_cards.append(MemoryCard.model_validate_json(line))
+
+ # Load embeddings
+ if os.path.exists(self._memory_embeddings_path):
+ self._memory_embeddings = np.load(self._memory_embeddings_path)
+ else:
+ self._memory_embeddings = np.zeros((len(self._memory_cards), 4096), dtype=np.float32)
+
+ # Load projection
+ if os.path.exists(self._item_projection_path):
+ proj_data = np.load(self._item_projection_path)
+ self._projection = ItemProjection(P=proj_data["P"], mean=proj_data["mean"])
+ self._item_vectors = proj_data["V"]
+ else:
+ self._projection = None
+ self._item_vectors = np.zeros((len(self._memory_cards), self._rl_cfg["item_dim"]), dtype=np.float32)
+
+ print(f"[PersonalizedLLM] Loaded {len(self._memory_cards)} memory cards.")
+
+ def _get_or_create_session(self, user_id: str) -> _SessionContext:
+ """Get or create session context for a user."""
+ if user_id not in self._sessions:
+ self._sessions[user_id] = _SessionContext(
+ session_state=OnlineSessionState(user_id=user_id),
+ turn_counter=0,
+ )
+ return self._sessions[user_id]
+
+ def _build_chat_turn(self, user_id: str, text: str, role: str, turn_id: int) -> ChatTurn:
+ """Build a ChatTurn object."""
+ return ChatTurn(
+ user_id=user_id,
+ session_id=f"eval_session_{user_id}",
+ turn_id=turn_id,
+ role=role,
+ text=text,
+ meta={"source": "eval"}
+ )
+
+ def _count_tokens(self, text: str) -> int:
+ """Estimate token count using the tokenizer."""
+ try:
+ # Use the chat model's tokenizer if available
+ if hasattr(self._chat_model, 'tokenizer'):
+ return len(self._chat_model.tokenizer.encode(text))
+ else:
+ # Rough estimate: ~4 chars per token
+ return len(text) // 4
+ except Exception:
+ return len(text) // 4
+
+ def _add_preferences_as_memory(
+ self,
+ prefs: PreferenceList,
+ query: str,
+ user_id: str,
+ turn_id: int,
+ ) -> List[Dict[str, Any]]:
+ """
+ Add extracted preferences as new memory cards.
+ Returns list of preference dicts for debug info.
+ """
+ extracted = []
+
+ if not prefs.preferences or self._projection is None:
+ return extracted
+
+ # Compute embedding for the query
+ e_q = self._embed_model.encode([query], return_tensor=False)[0]
+ v_q = self._projection.transform_vector(np.array(e_q))
+
+ for pref in prefs.preferences:
+ note_text = f"When {pref.condition}, {pref.action}."
+
+ # Record for debug
+ extracted.append({
+ "condition": pref.condition,
+ "action": pref.action,
+ "confidence": pref.confidence,
+ })
+
+ # Deduplication check
+ is_duplicate = any(
+ card.user_id == user_id and card.note_text == note_text
+ for card in self._memory_cards
+ )
+
+ if is_duplicate:
+ continue
+
+ # Create new memory card
+ card = MemoryCard(
+ card_id=str(uuid.uuid4()),
+ user_id=user_id,
+ source_session_id=f"eval_session_{user_id}",
+ source_turn_ids=[turn_id],
+ raw_queries=[query],
+ preference_list=PreferenceList(preferences=[pref]),
+ note_text=note_text,
+ embedding_e=list(e_q),
+ kind="pref",
+ )
+
+ # Add to memory store
+ self._memory_cards.append(card)
+ self._memory_embeddings = np.vstack([self._memory_embeddings, np.array([e_q])])
+ self._item_vectors = np.vstack([self._item_vectors, np.array([v_q])])
+
+ return extracted
+
+ # =========================================================================
+ # Public Interface
+ # =========================================================================
+
+ def chat(self, user_id: str, query: str) -> AssistantResponse:
+ """
+ Main online chat interface.
+
+ Args:
+ user_id: Unique identifier for the user.
+ query: Current user query/message.
+
+ Returns:
+ AssistantResponse containing the answer, usage stats, and debug info.
+
+ Notes:
+ - Internally manages user state, session history, memory retrieval
+ - After this call, you can call apply_feedback() with the turn's feedback
+ """
+ ctx = self._get_or_create_session(user_id)
+ session = ctx.session_state
+ user_state = self._user_store.get_state(user_id)
+
+ # Record user vector before for debug
+ z_long_before = user_state.z_long.copy().tolist()
+ z_short_before = user_state.z_short.copy().tolist()
+
+ # Compute query embedding
+ e_q_t = np.array(self._embed_model.encode([query], return_tensor=False)[0])
+
+ # Store pending RL update info from last turn (for apply_feedback)
+ if session.last_query is not None and self.enable_rl_updates:
+ ctx.pending_rl_update = {
+ "last_query": session.last_query,
+ "last_answer": session.last_answer,
+ "last_memories": session.last_memories,
+ "last_query_embedding": session.last_query_embedding,
+ "current_query_embedding": e_q_t,
+ "last_candidate_item_vectors": session.last_candidate_item_vectors,
+ "last_policy_probs": session.last_policy_probs,
+ "last_chosen_indices": session.last_chosen_indices,
+ }
+
+ # Add user turn to history
+ user_turn = self._build_chat_turn(user_id, query, "user", ctx.turn_counter)
+ session.history.append(user_turn)
+
+ # Vanilla mode: pure LLM without any memory or preference extraction
+ if self.mode == "vanilla":
+ # Skip preference extraction and memory retrieval entirely
+ extracted_prefs = []
+ candidates = []
+ cand_item_vecs = np.array([])
+ base_scores = np.array([])
+ chosen_indices = []
+ probs = np.array([])
+ memories_t = []
+ memory_notes = []
+ else:
+ # Extract preferences from conversation (if enabled)
+ extracted_prefs = []
+ if self.enable_preference_extraction:
+ prefs = self._extractor.extract_turn(session.history)
+ extracted_prefs = self._add_preferences_as_memory(
+ prefs, query, user_id, ctx.turn_counter
+ )
+
+ # Retrieve memories
+ # In "nopersonal" mode: deterministic retrieval (dense + rerank + topk), no policy/user vector
+ # In "full" mode: policy-based retrieval with user vector influence
+ if self.mode == "nopersonal":
+ candidates, cand_item_vecs, base_scores, chosen_indices, probs = retrieve_no_policy(
+ user_id=user_id,
+ query=query,
+ embed_model=self._embed_model,
+ reranker=self._reranker,
+ memory_cards=self._memory_cards,
+ memory_embeddings=self._memory_embeddings,
+ topk_dense=self._rl_cfg["dense_topk"],
+ topk_rerank=self._rl_cfg["rerank_topk"],
+ only_own_memories=self.only_own_memories,
+ )
+ else:
+ beta_long = self._rl_cfg["beta_long"]
+ beta_short = self._rl_cfg["beta_short"]
+ # eval_mode=True -> sample=False (greedy/deterministic)
+ # eval_mode=False -> sample=True (stochastic/exploration)
+ candidates, cand_item_vecs, base_scores, chosen_indices, probs = retrieve_with_policy(
+ user_id=user_id,
+ query=query,
+ embed_model=self._embed_model,
+ reranker=self._reranker,
+ memory_cards=self._memory_cards,
+ memory_embeddings=self._memory_embeddings,
+ user_store=self._user_store,
+ item_vectors=self._item_vectors,
+ topk_dense=self._rl_cfg["dense_topk"],
+ topk_rerank=self._rl_cfg["rerank_topk"],
+ beta_long=beta_long,
+ beta_short=beta_short,
+ tau=self._rl_cfg["tau"],
+ only_own_memories=self.only_own_memories,
+ sample=not self.eval_mode,
+ )
+
+ # Get selected memories
+ memories_t = [candidates[int(i)] for i in chosen_indices] if chosen_indices else []
+ memory_notes = [m.note_text for m in memories_t]
+
+ # Build prompt and count tokens
+ prompt_tokens = self._count_tokens(query)
+ for turn in session.history:
+ prompt_tokens += self._count_tokens(turn.text)
+ for note in memory_notes:
+ prompt_tokens += self._count_tokens(note)
+
+ # Generate answer
+ answer_t = self._chat_model.answer(
+ history=session.history,
+ memory_notes=memory_notes,
+ max_new_tokens=self._rl_cfg["max_new_tokens"],
+ )
+
+ completion_tokens = self._count_tokens(answer_t)
+
+ # Add assistant turn to history
+ assist_turn = self._build_chat_turn(user_id, answer_t, "assistant", ctx.turn_counter)
+ session.history.append(assist_turn)
+
+ # Update session state for next turn
+ session.last_query = query
+ session.last_answer = answer_t
+ session.last_memories = memories_t
+ session.last_query_embedding = e_q_t
+ session.last_candidate_item_vectors = cand_item_vecs
+ session.last_policy_probs = probs
+ session.last_chosen_indices = list(chosen_indices) if len(chosen_indices) > 0 else []
+
+ ctx.turn_counter += 1
+
+ # Build debug info
+ debug = DebugInfo(
+ selected_memory_ids=[m.card_id for m in memories_t],
+ selected_memory_notes=[m.note_text for m in memories_t],
+ selected_memory_scores=[float(probs[i]) if i < len(probs) else 0.0 for i in chosen_indices] if len(chosen_indices) > 0 else [],
+ user_vector_before=z_long_before + z_short_before, # Concatenated for simplicity
+ user_vector_after=user_state.z_long.tolist() + user_state.z_short.tolist(),
+ extracted_preferences=extracted_prefs,
+ extra={
+ "num_candidates": len(candidates),
+ "num_total_memories": len(self._memory_cards),
+ "z_long_norm": float(np.linalg.norm(user_state.z_long)),
+ "z_short_norm": float(np.linalg.norm(user_state.z_short)),
+ }
+ )
+
+ # Build usage stats
+ usage = UsageStats(
+ prompt_tokens=prompt_tokens,
+ completion_tokens=completion_tokens,
+ total_tokens=prompt_tokens + completion_tokens,
+ model=self._llm_name,
+ )
+
+ return AssistantResponse(
+ answer=answer_t,
+ usage=usage,
+ debug=debug,
+ )
+
+ def reset_session(self, user_id: str) -> None:
+ """
+ Reset session for a user (new chat window).
+
+ This clears:
+ - Session conversation history
+ - Short-term user vector (z_short)
+ - Pending RL update info
+
+ This preserves:
+ - Long-term user vector (z_long)
+ - User's memory cards
+
+ Args:
+ user_id: The user whose session to reset.
+ """
+ # Clear session context
+ if user_id in self._sessions:
+ del self._sessions[user_id]
+
+ # Create fresh session
+ self._sessions[user_id] = _SessionContext(
+ session_state=OnlineSessionState(user_id=user_id),
+ turn_counter=0,
+ )
+
+ # Reset short-term vector but keep long-term
+ user_state = self._user_store.get_state(user_id)
+ user_state.z_short = np.zeros(self._rl_cfg["item_dim"], dtype=np.float32)
+ self._user_store.save_state(user_state)
+
+ def reset_user(self, user_id: str) -> None:
+ """
+ Completely reset a user (new "life").
+
+ This clears:
+ - Long-term user vector (z_long)
+ - Short-term user vector (z_short)
+ - User's memory cards
+ - Session history
+ - All cached state
+
+ Args:
+ user_id: The user to reset.
+ """
+ # Clear session
+ if user_id in self._sessions:
+ del self._sessions[user_id]
+
+ # Reset user state vectors
+ user_state = self._user_store.get_state(user_id)
+ user_state.z_long = self._user_store.global_init_z.copy()
+ user_state.z_short = np.zeros(self._rl_cfg["item_dim"], dtype=np.float32)
+ user_state.reward_ma = 0.0
+ self._user_store.save_state(user_state)
+
+ # Find indices to KEEP (cards NOT belonging to this user)
+ # Must do this BEFORE modifying _memory_cards
+ keep_indices = [
+ i for i, card in enumerate(self._memory_cards)
+ if card.user_id != user_id
+ ]
+
+ # Filter memory cards
+ self._memory_cards = [self._memory_cards[i] for i in keep_indices]
+
+ # Filter embeddings and item vectors to match
+ if len(keep_indices) > 0 and len(self._memory_embeddings) > 0:
+ self._memory_embeddings = self._memory_embeddings[keep_indices]
+ self._item_vectors = self._item_vectors[keep_indices]
+ else:
+ # No cards left or no embeddings
+ embed_dim = self._memory_embeddings.shape[1] if len(self._memory_embeddings) > 0 else 4096
+ self._memory_embeddings = np.zeros((0, embed_dim), dtype=np.float32)
+ self._item_vectors = np.zeros((0, self._rl_cfg["item_dim"]), dtype=np.float32)
+
+ def apply_feedback(self, feedback: Feedback) -> None:
+ """
+ Apply feedback from user simulator or judge.
+
+ This performs the REINFORCE update to user vectors based on
+ the reward signal from the previous turn.
+
+ Args:
+ feedback: Feedback object containing reward, gating, and metadata.
+
+ Notes:
+ - Should be called AFTER chat() but BEFORE the next chat() call
+ - Uses the stored context from the previous turn
+ - If enable_rl_updates is False, this is a no-op (logging only)
+ - If mode is "nopersonal", this is a no-op (baseline comparison)
+ """
+ if not self.enable_rl_updates:
+ return
+
+ # In "nopersonal" or "vanilla" mode, skip RL updates entirely (baseline)
+ if self.mode in ("nopersonal", "vanilla"):
+ return
+
+ user_id = feedback.user_id
+ ctx = self._sessions.get(user_id)
+
+ if ctx is None or ctx.pending_rl_update is None:
+ return
+
+ pending = ctx.pending_rl_update
+ user_state = self._user_store.get_state(user_id)
+
+ # Check if we have the necessary data for RL update
+ if (pending.get("last_candidate_item_vectors") is not None and
+ pending.get("last_policy_probs") is not None and
+ pending.get("last_chosen_indices") is not None and
+ len(pending["last_chosen_indices"]) > 0):
+
+ # Extract chosen vectors
+ chosen_indices = pending["last_chosen_indices"]
+ candidate_vectors = pending["last_candidate_item_vectors"]
+
+ if len(candidate_vectors) > 0:
+ # REINFORCE expects:
+ # - item_vectors: ALL candidate vectors [K, k]
+ # - chosen_indices: indices into those candidates
+ # - policy_probs: probabilities over all K candidates [K]
+ updated = reinforce_update_user_state(
+ user_state=user_state,
+ item_vectors=candidate_vectors, # All candidates, not just chosen
+ chosen_indices=chosen_indices, # Original indices into candidates
+ policy_probs=pending["last_policy_probs"],
+ reward_hat=feedback.reward,
+ gating=feedback.gating,
+ tau=self._rl_cfg["tau"],
+ eta_long=self._rl_cfg["eta_long"],
+ eta_short=self._rl_cfg["eta_short"],
+ ema_alpha=self._rl_cfg["ema_alpha"],
+ short_decay=self._rl_cfg["short_decay"],
+ )
+
+ if updated:
+ self._user_store.save_state(user_state)
+
+ # Clear pending update
+ ctx.pending_rl_update = None
+
+ def get_user_state_summary(self, user_id: str) -> Dict[str, Any]:
+ """
+ Get a summary of the user's current state (for debugging/analysis).
+
+ Args:
+ user_id: The user to query.
+
+ Returns:
+ Dictionary with user state information.
+ """
+ user_state = self._user_store.get_state(user_id)
+ ctx = self._sessions.get(user_id)
+
+ user_memory_count = sum(
+ 1 for card in self._memory_cards if card.user_id == user_id
+ )
+
+ return {
+ "user_id": user_id,
+ "z_long_norm": float(np.linalg.norm(user_state.z_long)),
+ "z_short_norm": float(np.linalg.norm(user_state.z_short)),
+ "reward_ma": user_state.reward_ma,
+ "session_history_length": len(ctx.session_state.history) if ctx else 0,
+ "turn_counter": ctx.turn_counter if ctx else 0,
+ "user_memory_count": user_memory_count,
+ "total_memory_count": len(self._memory_cards),
+ }
+
+ def persist(self) -> None:
+ """
+ Persist all state to disk.
+
+ Call this at the end of an evaluation run to save:
+ - User state vectors
+ - Memory cards
+ """
+ # Save user store
+ self._user_store.persist()
+
+ # Save memory cards
+ with open(self._memory_cards_path, "w", encoding="utf-8") as f:
+ for card in self._memory_cards:
+ f.write(card.model_dump_json() + "\n")
+
+ # Save embeddings
+ np.save(self._memory_embeddings_path, self._memory_embeddings)
+
+ # Save item projection with updated vectors
+ if self._projection is not None:
+ np.savez(
+ self._item_projection_path,
+ P=self._projection.P,
+ mean=self._projection.mean,
+ V=self._item_vectors,
+ )
+
+ print("[PersonalizedLLM] State persisted to disk.")
+
+
+# =============================================================================
+# Convenience Factory
+# =============================================================================
+
+def create_personalized_llm(
+ config_path: Optional[str] = None,
+ **kwargs
+) -> PersonalizedLLM:
+ """
+ Factory function to create a PersonalizedLLM instance.
+
+ Args:
+ config_path: Optional path to configuration file.
+ **kwargs: Additional arguments passed to PersonalizedLLM constructor.
+
+ Returns:
+ Configured PersonalizedLLM instance.
+ """
+ return PersonalizedLLM(config_path=config_path, **kwargs)
+
generated by cgit v1.2.3 (git 2.25.1) at 2026-01-27 04:15:18 +0000