Upload handler.py

handler.py

@@ -21,9 +21,8 @@ import importlib.util
 from dataclasses import dataclass
 from pathlib import Path
 import hashlib
-import random
 import re
-from typing import Any, Mapping, MutableMapping, Optional, Sequence
+from typing import Any, Mapping, MutableMapping, Optional, Sequence, Tuple
 
 import numpy as np
 
@@ -380,61 +379,133 @@ class EndpointHandler:
 
         return self._coerce_array(value, node=node)
 
-    def _encode_tokens(self, text: str) -> tuple[np.ndarray, list[int]]:
-        token_ids = self._tokenizer.encode(text)
+    @staticmethod
+    def _candidate_seed(psi: np.ndarray) -> int:
+        digest = hashlib.sha1(psi.tobytes()).digest()
+        return int.from_bytes(digest[:4], "little", signed=False)
+
+    def _token_array_from_ids(self, token_ids: Sequence[int]) -> np.ndarray:
+        ids = list(token_ids)
         if self._token_sequence_length <= 0:
-            array = np.asarray([token_ids], dtype=self._token_dtype)
-            return array, token_ids
+            return np.asarray([ids], dtype=self._token_dtype)
 
-        length = min(len(token_ids), self._token_sequence_length)
         padded = np.full(
             (1, self._token_sequence_length),
             fill_value=self._tokenizer.pad_token_id,
             dtype=self._token_dtype,
         )
-        padded[0, :length] = np.asarray(token_ids[:length], dtype=self._token_dtype)
-        return padded, token_ids[:length]
+        length = min(len(ids), self._token_sequence_length)
+        if length > 0:
+            padded[0, :length] = np.asarray(ids[:length], dtype=self._token_dtype)
+        return padded
+
+    def _run_candidate(self, base_feed: Mapping[str, np.ndarray], tokens: Sequence[int]) -> list[tuple[Any, np.ndarray]]:
+        feed = {
+            name: (value.copy() if isinstance(value, np.ndarray) else value)
+            for name, value in base_feed.items()
+        }
+        if self.tokens_input is not None:
+            feed[self.tokens_input] = self._token_array_from_ids(tokens)
+        outputs = self.session.run(None, feed)
+        return list(zip(self.io.outputs, outputs))
 
     @staticmethod
-    def _candidate_seed(psi: np.ndarray) -> int:
-        digest = hashlib.sha1(psi.tobytes()).digest()
-        return int.from_bytes(digest[:4], "little", signed=False)
+    def _extract_logits(outputs: Sequence[tuple[Any, np.ndarray]]) -> Optional[np.ndarray]:
+        for node, value in outputs:
+            if getattr(node, "name", "").lower() == "logits":
+                return np.asarray(value, dtype=np.float32)
+        if outputs:
+            return np.asarray(outputs[0][1], dtype=np.float32)
+        return None
 
-    def _build_candidates(
+    @staticmethod
+    def _sample_next_token(
+        logits: np.ndarray,
+        decoding: _DecodingParams,
+        rng: np.random.Generator,
+    ) -> int:
+        vector = np.asarray(logits, dtype=np.float64).reshape(-1)
+        temperature = max(float(decoding.temperature), 1e-5)
+        top_p = float(decoding.top_p)
+
+        if temperature <= 1e-5 or not np.isfinite(vector).any():
+            return int(int(np.argmax(vector)))
+
+        stabilized = vector / temperature
+        stabilized -= np.max(stabilized)
+        probs = np.exp(stabilized)
+        probs = np.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+        total = probs.sum()
+        if total <= 0.0:
+            return int(np.argmax(vector))
+        probs /= total
+
+        if top_p <= 0.0:
+            return int(np.argmax(probs))
+
+        if 0.0 < top_p < 1.0:
+            sorted_indices = np.argsort(-probs)
+            sorted_probs = probs[sorted_indices]
+            cumulative = np.cumsum(sorted_probs)
+            mask = cumulative <= top_p
+            if mask.size > 0:
+                mask[0] = True
+            filtered_indices = sorted_indices[mask]
+            filtered_probs = sorted_probs[mask]
+            filtered_total = filtered_probs.sum()
+            if filtered_total <= 0.0:
+                filtered_indices = sorted_indices
+                filtered_probs = sorted_probs
+                filtered_total = filtered_probs.sum()
+            filtered_probs = filtered_probs / filtered_total
+            choice = rng.choice(len(filtered_indices), p=filtered_probs)
+            return int(filtered_indices[int(choice)])
+
+        choice = rng.choice(len(probs), p=probs)
+        return int(choice)
+
+    def _generate_sequence(
         self,
-        psi_vector: np.ndarray,
+        base_feed: Mapping[str, np.ndarray],
         *,
-        user_prompt: str | None,
-        system_prompt: str | None,
-        constraints: Mapping[str, Any] | None,
-    ) -> tuple[list[str], str, list[str]]:
-        descriptors = _summarise_intent(psi_vector)
-        summary = ", ".join(descriptors) if descriptors else "balanced intent"
-        observations = [
-            f"Interpretation: the symbolic intent emphasises {summary}.",
-            f"Symbolic synopsis → {summary}.",
-        ]
-        if user_prompt:
-            observations.append(f"{user_prompt.strip()}\nInsight: {summary}.")
-        if system_prompt:
-            observations.append(f"{system_prompt.strip()}\nDirective: honour {summary}.")
-        if constraints:
-            formatted = ", ".join(f"{key}={value}" for key, value in constraints.items())
-            observations.append(f"Constraints observed: {formatted}.")
-
-        seed = self._candidate_seed(psi_vector.astype(np.float32, copy=False))
-        rng = random.Random(seed)
-        rng.shuffle(observations)
-        if not observations:
-            observations = [f"Symbolic synopsis → {summary}."]
-        return observations, summary, descriptors
-
-    def _run_candidate(self, base_feed: Mapping[str, np.ndarray], tokens: np.ndarray) -> list[tuple[Any, np.ndarray]]:
-        feed = {name: value for name, value in base_feed.items()}
-        if self.tokens_input is not None:
-            feed[self.tokens_input] = tokens
-        outputs = self.session.run(None, feed)
-        return list(zip(self.io.outputs, outputs))
+        decoding: _DecodingParams,
+        seed: int,
+    ) -> Optional[Tuple[str, list[int], float, list[tuple[Any, np.ndarray]], int]]:
+        if self.tokens_input is None:
+            return None
+
+        rng = np.random.default_rng(seed)
+        token_ids: list[int] = [self._tokenizer.bos_token_id]
+        quality = float("-inf")
+        formatted_outputs: list[tuple[Any, np.ndarray]] | None = None
+        steps = 0
+
+        max_steps = max(decoding.max_new_tokens, 1)
+        for _ in range(max_steps):
+            outputs = self._run_candidate(base_feed, token_ids)
+            logits = self._extract_logits(outputs)
+            if logits is None:
+                break
+            last_index = min(len(token_ids) - 1, logits.shape[1] - 1)
+            next_logits = logits[0, last_index]
+            next_token = self._sample_next_token(next_logits, decoding, rng)
+            token_ids.append(int(next_token))
+            steps += 1
+
+            outputs = self._run_candidate(base_feed, token_ids)
+            formatted_outputs = outputs
+            quality = self._extract_q_hat(outputs)
+
+            if token_ids[-1] == self._tokenizer.eos_token_id:
+                break
+            if self._token_sequence_length > 0 and len(token_ids) >= self._token_sequence_length:
+                break
+
+        if formatted_outputs is None:
+            return None
+
+        text = self._tokenizer.decode(token_ids)
+        return text, token_ids, float(quality), formatted_outputs, steps
 
     @staticmethod
     def _extract_q_hat(outputs: Sequence[tuple[Any, np.ndarray]]) -> float:
@@ -470,66 +541,72 @@ class EndpointHandler:
         system_prompt = payload.get("system_prompt")
         user_prompt = payload.get("user_prompt")
 
-        candidates, summary, descriptors = self._build_candidates(
-            psi_vector,
-            user_prompt=user_prompt if isinstance(user_prompt, str) else None,
-            system_prompt=system_prompt if isinstance(system_prompt, str) else None,
-            constraints=state_constraints,
-        )
-
-        best_text: str | None = None
-        best_tokens: list[int] = []
-        best_outputs: list[tuple[Any, np.ndarray]] | None = None
-        best_quality = float("-inf")
+        descriptors = _summarise_intent(psi_vector)
+        summary = ", ".join(descriptors) if descriptors else "balanced intent"
 
-        limit = min(len(candidates), max(decoding.beam_size, 1))
-        for candidate in candidates[:limit]:
-            if self.tokens_input is None:
-                break
-            token_array, token_ids = self._encode_tokens(candidate)
-            outputs = self._run_candidate(feed, token_array)
-            quality = self._extract_q_hat(outputs)
-            if quality > best_quality:
-                best_quality = quality
-                best_text = candidate
-                best_tokens = token_ids
-                best_outputs = outputs
-            if quality >= decoding.stop_quality:
-                break
+        best_candidate: Optional[Tuple[str, list[int], float, list[tuple[Any, np.ndarray]], int]] = None
+        seeds: list[int] = []
 
-        if best_outputs is None:
-            # Fall back to a single pass using the prepared feed.
+        if self.tokens_input is not None:
+            beams = max(decoding.beam_size, 1)
+            base_seed = self._candidate_seed(psi_vector)
+            for beam_idx in range(beams):
+                seed = base_seed + beam_idx
+                seeds.append(seed)
+                candidate = self._generate_sequence(
+                    feed,
+                    decoding=decoding,
+                    seed=seed,
+                )
+                if candidate is None:
+                    continue
+                text, token_ids, quality, outputs, steps = candidate
+                if (
+                    best_candidate is None
+                    or quality > best_candidate[2]
+                ):
+                    best_candidate = candidate
+                if quality >= decoding.stop_quality:
+                    break
+
+        if best_candidate is None:
             outputs = self.session.run(None, feed)
-            best_outputs = list(zip(self.io.outputs, outputs))
-            if best_text is None:
-                best_text = f"Symbolic synopsis → {summary}."
-            if best_quality == float("-inf"):
-                best_quality = self._extract_q_hat(best_outputs)
+            formatted_outputs = list(zip(self.io.outputs, outputs))
+            quality = self._extract_q_hat(formatted_outputs)
+            text = f"Symbolic synopsis → {summary}."
+            token_ids: list[int] = []
+            steps = 0
+        else:
+            text, token_ids, quality, formatted_outputs, steps = best_candidate
 
         formatted = {
             node.name: self._format_output(node.name, value)
-            for node, value in best_outputs
+            for node, value in formatted_outputs
         }
 
-        if not np.isfinite(best_quality):
-            best_quality = 0.0
-        best_quality = float(best_quality)
-        if best_text is None:
-            best_text = f"Symbolic synopsis → {summary}."
+        if not np.isfinite(quality):
+            quality = 0.0
+        quality = float(quality)
+
+        metadata = {
+            "summary": summary,
+            "descriptors": descriptors,
+            "constraints": state_constraints or {},
+            "decoding": decoding.to_dict(),
+            "seeds": seeds,
+            "steps": steps,
+            "system_prompt": system_prompt if isinstance(system_prompt, str) else None,
+            "user_prompt": user_prompt if isinstance(user_prompt, str) else None,
+        }
 
         response = {
-            "text": best_text,
-            "tokens": best_tokens,
-            "quality": best_quality,
-            "q_hat": best_quality,
+            "text": text,
+            "tokens": token_ids,
+            "quality": quality,
+            "q_hat": quality,
             "provider": _DEFAULT_PROVIDER,
             "model": _DEFAULT_MODEL,
-            "metadata": {
-                "summary": summary,
-                "descriptors": descriptors,
-                "constraints": state_constraints or {},
-                "decoding": decoding.to_dict(),
-            },
+            "metadata": metadata,
         }
         response.update(formatted)
         return response