H2O v1.4.2 serving files (weights unchanged)

.h2o_revision

@@ -1 +1 @@
-d7f1b05aa0b4e789b3896116f309ef842440e260
+708e93926a034195d908e033fcb796c55398a6f5

h2o_lightning_shim.py

@@ -94,6 +94,15 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 
 SHIM_VERSION = "h2o-lightning-shim"
 PRIMITIVES = ("choice", "noul", "score")
+# THE DI PROFILE'S CONDITIONAL RE-READ (off unless SHIM_REREAD_BELOW is set; serve.sh PROFILE=di sets 0.8). A text
+# question whose served top probability is under the threshold (a yes/no answer's max(P, 1 - P) after the floor; PROFILE=di
+# pins the floor at 0.801, so there yes/no questions are read once; with the floor off they would be eligible) is read once more in the same attention mode with the user turn written
+# twice, the second copy after REREAD_SEP, and is answered with the 2:1 first-weighted mean of the two reads'
+# probability vectors (REREAD_FIRST_WEIGHT; chosen on held-out data after the plain-mean rule passed). The re-read is
+# skipped (the first answer stands) when the question's own text was cut to fit, the doubled prompt would not fit the
+# engine's context, or the engine refuses it: a request never fails because of the re-read. Images are read once.
+REREAD_SEP = "\n\nRead the same input again before answering:\n"
+REREAD_FIRST_WEIGHT = 2.0 / 3.0
 EFFORT_TIERS = ("fast", "balanced", "thorough")
 # vLLM 0.30 accepts at most 128 `logprob_token_ids` per request; a question with more options is read in chunks
 LOGPROB_CHUNK = int(os.environ.get("SHIM_LOGPROB_CHUNK", "128"))
@@ -765,6 +774,13 @@ def format_answer(typ, names, descs, logprobs, temperature, noul_floor, threshol
     }
 
 
+
+def served_top(answer):
+    """An answer's top probability as served: a yes/no answer's max(P, 1 - P) (after the floor), else its top option."""
+    if answer["type"] == "noul":
+        return max(answer["noul"], 1.0 - answer["noul"])
+    return max(answer["probabilities"].values())
+
 MODES = ("causal", "bidirectional")
 MODE_ALIASES = {"causal": "causal", "one-way": "causal", "one_way": "causal", "oneway": "causal",
                 "bidirectional": "bidirectional", "two-way": "bidirectional", "two_way": "bidirectional",
@@ -958,6 +974,10 @@ class Contract:
         self.temperature_source = "SHIM_TEMPERATURE" if t_env else "serve_config.json"
         self.noul_floor = float(f_env) if f_env else float(cfg["noul_floor"])
         self.noul_floor_source = "SHIM_NOUL_FLOOR" if f_env else "serve_config.json"
+        rb = env.get("SHIM_REREAD_BELOW", "").strip()
+        self.reread_below = float(rb) if rb else None
+        if self.reread_below is not None and not 0.0 < self.reread_below <= 1.0:
+            raise ValueError(f"SHIM_REREAD_BELOW={rb!r}: a probability in (0, 1], or unset (read once)")
         if not all(0.05 <= t <= 20.0 for t in self.temperature_by_type.values()):
             raise ValueError(f"temperature {self.temperature_by_type} is outside [0.05, 20]")
         if self.noul_floor and not 0.5 < self.noul_floor < 1.0:
@@ -1022,18 +1042,20 @@ class Contract:
         tw = cfg.get("two_way")
         self.two_way = TwoWay(tw, self, cfg.get("_config_dir"), env) if isinstance(tw, dict) else None
 
-    def question_parts(self, state_text, q, qid="q"):
-        """(user text, option names, descriptions, labels) for one question."""
+    def question_parts(self, state_text, q, qid="q", twice=False):
+        """(user text, option names, descriptions, labels) for one question; `twice`: the user turn written again after
+        REREAD_SEP (the DI profile's re-read)."""
         names, descs = options_of(q, len(self.labels))
         if len(names) > len(self.labels):
             raise Unprocessable(f"question {qid}: {len(names)} options exceeds {len(self.labels)}")
         labels = self.labels[:len(names)]
         instr = render_instr(q.get("instructions", "")).lstrip() or self.default_instructions
-        return user_turn(state_text, instr, names, descs, labels, self.layout), names, descs, labels
+        user = user_turn(state_text, instr, names, descs, labels, self.layout)
+        return (user + REREAD_SEP + user if twice else user), names, descs, labels
 
-    def question_prompt(self, state_text, q, qid="q"):
+    def question_prompt(self, state_text, q, qid="q", twice=False):
         """(prompt text, option names, descriptions, labels) for one question."""
-        user, names, descs, labels = self.question_parts(state_text, q, qid)
+        user, names, descs, labels = self.question_parts(state_text, q, qid, twice)
         return render_chat(self.system, user, self.prefill, self.chat), names, descs, labels
 
     def image_prompt(self, state_text, q, image_urls, qid="q"):
@@ -1400,9 +1422,16 @@ class Shim:
         return self.c.two_way.default if self.v is not None else "bidirectional"
 
     def mode(self, body):
-        """The attention mode of a request: its `attention_mode`, else the deployment's default. A mode whose engine
-        is not running is a 400 (never a silent answer in the other mode)."""
-        asked = body.get("attention_mode") if isinstance(body, dict) else None
+        """The attention mode of a request: its `attention_mode` (top level, or `kwargs.attention_mode`, the form the
+        hosted gateway and the 31B package take; given in both with different values is a 400, as the gateway does),
+        else the deployment's default. A mode whose engine is not running is a 400 (never a silent answer in the other
+        mode)."""
+        asked = None
+        if isinstance(body, dict):
+            kw = body.get("kwargs") if isinstance(body.get("kwargs"), dict) else {}
+            if "attention_mode" in body and "attention_mode" in kw and body["attention_mode"] != kw["attention_mode"]:
+                raise BadRequest("attention_mode given twice with different values (top level and kwargs)")
+            asked = body.get("attention_mode", kw.get("attention_mode"))
         m = _mode_of(asked, self.default_mode())
         if self.two_way_flag:
             return m                      # one engine serves both through the per-request flag
@@ -1907,17 +1936,102 @@ class Shim:
             if not all(isinstance(t, list) for t in toks):
                 toks = [self.t.tokenize(w[4]) for w in work]
             shared = common_prefix_len(toks)
-        return self.respond(work, reads, shared, asked, t0, self.temperatures(mode), full_input=full, mode=mode,
-                            notes=notes)
+        if self.c.reread_below is None:
+            return self.respond(work, reads, shared, asked, t0, self.temperatures(mode), full_input=full, mode=mode,
+                                notes=notes)
+        rereads, extra = self.reread(work, reads, questions, fits, read, mode)
+        out = self.respond(work, reads, shared, asked, t0, self.temperatures(mode), full_input=full, mode=mode,
+                           notes=notes)
+        out["usage"]["input_tokens"] += extra
+        out["usage"]["processed_input_tokens"] += extra
+        out["effort_used"]["reread"] = {"below": self.c.reread_below, "first_weight": round(REREAD_FIRST_WEIGHT, 4), "questions": rereads}
+        return out
+
+    def reread(self, work, reads, questions, fits, read, mode):
+        """The DI profile's conditional re-read (REREAD_SEP): each text question whose SERVED first-read top
+        probability (the measured rule: a choice's or score's top option; a yes/no answer's max(P, 1 - P) after the
+        floor; PROFILE=di pins the floor at 0.801, so there yes/no questions are read once; with the floor off they
+        would be eligible) is under c.reread_below is
+        read again, in the same mode, with its user turn written twice, and is answered with the 2:1 first-weighted mean
+        of the two reads' probability vectors at the mode's temperature (REREAD_FIRST_WEIGHT). `reads` takes, in place
+        of the first read's log-probabilities, T * log(mean), whose softmax at T is that mean. Returns (questions
+        re-read, the re-reads' prompt tokens). A question whose own text was cut to fit, a doubled prompt over the
+        engine's context, or one the engine refuses keeps the first read: a request never fails because of the re-read."""
+        t_by, _it, floor = self.temperatures(mode)
+        todo = []
+        for i, ((qid, typ, names, descs, _text, labels, thr), (lp, _u), (state_txt, user)) in enumerate(zip(work, reads, fits)):
+            t = t_by[typ] if isinstance(t_by, dict) else t_by
+            if user is not None or served_top(format_answer(typ, names, descs, lp, t, floor, thr)) >= self.c.reread_below:
+                continue
+            todo.append((i, t, self.c.question_prompt(state_txt, questions[qid], qid, twice=True)[0], labels))
+        if not todo:
+            return 0, 0
+        if isinstance(self.max_model_len, int):
+            todo = [x for x in todo if len(self.t.tokenize(x[2])) + 8 <= self.max_model_len]
+
+        def again(text, labels):
+            try:
+                return read(text, labels, None)
+            except (UpstreamError, Unprocessable):
+                return None
+        futs = [(i, t, self.pool.submit(again, text, labels)) for i, t, text, labels in todo]
+        n, extra = 0, 0
+        for i, t, f in futs:
+            r = f.result()
+            if r is None:
+                continue
+            p1, p2 = probabilities(reads[i][0], t), probabilities(r[0], t)
+            reads[i] = ([t * math.log(max(REREAD_FIRST_WEIGHT * a + (1.0 - REREAD_FIRST_WEIGHT) * b, 1e-300))
+                         for a, b in zip(p1, p2)], reads[i][1])
+            extra += int(r[1].get("prompt_tokens") or 0)
+            n += 1
+        return n, extra
 
     def fit_request(self, raw_state, full_txt, questions, n_images=0):
         """[(record text, user-turn override) per question], notes: selector headtail cuts the record once (v1.4.1);
         selector v2 fits each question to the context (fit_states)."""
         if self.c.selector == "headtail":
             st = self.truncate_state(full_txt)
-            return [(st, None)] * len(questions), (["head+tail (selector headtail)"] if st != full_txt else [])
+            notes = ["head+tail (selector headtail)"] if st != full_txt else []
+            return self.fit_headtail(full_txt, st, questions, n_images, notes), notes
         return self.fit_states(raw_state, full_txt, questions, n_images)
 
+    def fit_headtail(self, full_txt, st, questions, n_images, notes):
+        """Selector headtail's per-question CONTEXT CLAMP: a question whose prompt with the head+tail record `st` fits
+        vLLM's context reads it as before; one that does not (SHIM_MAX_STATE_TOKENS near the context with a long
+        question) gets the record cut by head and tail to the room its own prompt leaves, the same per-question budget
+        selector v2 uses (max_model_len less the prompt without the record, the images' bound and CONTEXT_SLACK); a
+        question that does not fit even without the record is answered without it (fit_question). Never a 422 for
+        length."""
+        c, C = self.c, self.max_model_len
+        img = n_images * self.image_tokens_bound()
+        limit = (C - CONTEXT_SLACK) if isinstance(C, int) and C > 0 else None
+        out = []
+        for qid, q in questions.items():
+            text = c.question_prompt(st, q, qid)[0]
+            if limit is None or len(text.encode("utf-8", "ignore")) + img <= limit or len(self.t.tokenize(text)) + img <= limit:
+                out.append((st, None))
+                continue
+            room = limit - len(self.t.tokenize(c.question_prompt("", q, qid)[0])) - img
+            if room < MIN_STATE_BUDGET:
+                notes.append("question over the context; record left out")
+                out.append(("", self.fit_question(q, qid, limit, img, notes)))
+                continue
+            budget = min(room, c.max_state_tokens)
+            cut = self.cut_text(full_txt, budget)
+            for _ in range(3):                 # the exact count (the record's seams with the prompt can add a token)
+                n = len(self.t.tokenize(c.question_prompt(cut, q, qid)[0])) + img
+                if n <= limit:
+                    break
+                budget -= (n - limit) + 8
+                if budget < MIN_STATE_BUDGET:
+                    cut = ""
+                    break
+                cut = self.cut_text(full_txt, budget)
+            notes.append("head+tail cut to the context")
+            out.append((cut, None))
+        return out
+
     def full_input(self, prompts, image_tokens=0):
         """Input tokens as if the whole record had been read: the full prompts, their shared head once, plus the
         image tokens (BILLING = THE FULL INPUT: a cut record gains speed, not a lower count)."""
@@ -2036,6 +2150,7 @@ class Shim:
                 "noul_floor": {"value": c.noul_floor, "source": c.noul_floor_source},
                 "temperature_by_type": dict(c.temperature_by_type),
                 "temperature_source": c.temperature_source, "max_options": len(c.labels),
+                **({"reread_below": c.reread_below} if c.reread_below is not None else {}),
                 "vllm_max_model_len": self.max_model_len,
                 **({"images": {"max_images": c.max_images, "max_pixels": c.max_pixels or None,
                                "temperature": c.image_temperature, "verified": bool(self.chat_ok)}}
@@ -2225,11 +2340,14 @@ def make_handler(shim):
 
         def _send(self, code, obj):
             payload = json.dumps(obj).encode()
-            self.send_response(code)
-            self.send_header("Content-Type", "application/json")
-            self.send_header("Content-Length", str(len(payload)))
-            self.end_headers()
-            self.wfile.write(payload)
+            try:
+                self.send_response(code)
+                self.send_header("Content-Type", "application/json")
+                self.send_header("Content-Length", str(len(payload)))
+                self.end_headers()
+                self.wfile.write(payload)
+            except (BrokenPipeError, ConnectionResetError):   # the client left first (its own deadline): nobody to answer
+                self.close_connection = True
 
         def do_GET(self):
             if self.path.startswith("/health"):

serve_config.json

@@ -16,7 +16,7 @@
   "noul": 0.8,
   "score": 0.65
  },
- "noul_floor": 0.801,
+ "noul_floor": 0,
  "prompt": {
   "chat_template": true,
   "enable_thinking": false,
@@ -286,7 +286,7 @@
  "max_options": 255,
  "image": {
   "max_images": 4,
-  "max_pixels": 1048576,
+  "max_pixels": 1638400,
   "temperature": 1.0
  },
  "two_way": {
@@ -297,7 +297,7 @@
    "noul": 0.8,
    "score": 0.65
   },
-  "noul_floor": 0.801,
+  "noul_floor": 0,
   "image_temperature": 1.0,
   "probe_tolerance": 0.05,
   "probe_reference": [