Upload handler.py with huggingface_hub

handler.py

@@ -1,28 +1,212 @@
-"""Minimal handler for debugging."""
+"""Custom handler for HF Inference Endpoints.
+
+Loads the pre-merged Gemma 2B IT model and exposes the same inference
+logic as server.py — Gemma chat template, INSTRUCTION prompt, and
+generation parameters.
+
+Supports MC Dropout for Bayesian uncertainty estimation. Since the model
+is pre-merged (no LoRA dropout layers), we inject DropoutWrapper modules
+around attention projection layers at init time.
+"""
+
+import os
+import sys
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# On HF Inference Endpoints, handler.py and mc_dropout.py are in the model
+# repo root. Ensure that directory is on sys.path for all imports.
+_handler_dir = os.path.dirname(os.path.abspath(__file__))
+if _handler_dir not in sys.path:
+    sys.path.insert(0, _handler_dir)
+
+# MC Dropout import — gracefully degrade if unavailable
+_MC_AVAILABLE = False
+try:
+    from mc_dropout import (
+        aggregate_beliefs,
+        disable_mc_dropout,
+        enable_mc_dropout,
+        inject_dropout,
+    )
+    _MC_AVAILABLE = True
+except Exception as _mc_err:
+    # Try local development path
+    try:
+        sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+        from mc_dropout import (
+            aggregate_beliefs,
+            disable_mc_dropout,
+            enable_mc_dropout,
+            inject_dropout,
+        )
+        _MC_AVAILABLE = True
+    except Exception:
+        print(f"WARNING: MC Dropout not available: {_mc_err}")
+        aggregate_beliefs = None
+        disable_mc_dropout = None
+        enable_mc_dropout = None
+        inject_dropout = None
+
+INSTRUCTION = """\
+Narrative format: {timestamp} {hash8} {refs} {actors} :: {subject} | {sym}{file}+N/-N@{funcs} ...
+Symbols: + added, ~ modified, - deleted, > renamed, = copied
+Roles: (a+c) author+committer, (a) author, (c) committer
+Funcs: @{name1,name2} for modified functions/classes
+Merges: {timestamp} {hash8} MERGE {merged}→{main} :: {subject}
+Renames: >{old}→{new}(N%)+N/-N  Large commits: ...+Nmore when >50 files
+
+Read the repository events below. Identify recurring patterns in workflow, code ownership, commit discipline, and architecture.
+
+For each pattern found, output:
+Belief: <clear statement about a development practice or pattern>
+Evidence:
+- Commit: <hash(es) that support this belief>
+- Files: <file paths involved>
+- Pattern: <what pattern was observed and why it matters>
+Confidence: high | medium | low"""
+
+CHUNK_SIZE = 4
+
 
 class EndpointHandler:
     def __init__(self, path):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         dtype = torch.float16 if self.device == "cuda" else torch.float32
+
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
             path, torch_dtype=dtype
         ).to(self.device)
         self.model.eval()
 
-    def __call__(self, data):
-        inputs = data.get("inputs", "test")
+        # Inject dropout wrappers for MC Dropout support on merged model
+        self.mc_dropout_available = False
+        if _MC_AVAILABLE and inject_dropout is not None:
+            mc_dropout_rate = float(os.environ.get("MC_DROPOUT_RATE", "0.1"))
+            n_injected = inject_dropout(self.model, dropout_rate=mc_dropout_rate)
+            print(f"Injected {n_injected} dropout wrappers (rate={mc_dropout_rate})")
+            self.mc_dropout_available = n_injected > 0
+        else:
+            print("MC Dropout not available — running in standard mode")
+
+    def _run_inference(self, text):
+        """Run inference on a single text chunk using Gemma chat template."""
+        user_content = f"{INSTRUCTION}\n\n{text}"
         prompt = (
             f"<start_of_turn>user\n"
-            f"{inputs}<end_of_turn>\n"
+            f"{user_content}<end_of_turn>\n"
             f"<start_of_turn>model\n"
         )
-        tok = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(self.device)
+
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048,
+        ).to(self.device)
+
         with torch.no_grad():
-            out = self.model.generate(**tok, max_new_tokens=100, pad_token_id=self.tokenizer.eos_token_id)
-        text = self.tokenizer.decode(out[0][tok["input_ids"].shape[1]:], skip_special_tokens=True).strip()
-        return {"generated_text": text}
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=400,
+                do_sample=False,
+                pad_token_id=self.tokenizer.eos_token_id,
+                eos_token_id=[self.tokenizer.eos_token_id, 107],
+            )
+
+        generated = self.tokenizer.decode(
+            outputs[0][inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+        ).strip()
+
+        return generated
+
+    @staticmethod
+    def _chunk_narrative(narrative, chunk_size=CHUNK_SIZE):
+        """Split narrative into non-overlapping chunks."""
+        lines = [
+            line.rstrip("\n") for line in narrative.splitlines() if line.strip()
+        ]
+        chunks = []
+        for i in range(0, len(lines), chunk_size):
+            chunk_lines = lines[i : i + chunk_size]
+            if chunk_lines:
+                chunks.append("\n".join(chunk_lines))
+        return chunks
+
+    def _run_mc_inference(self, text, n_passes):
+        """Run MC Dropout inference: N passes with dropout enabled."""
+        user_content = f"{INSTRUCTION}\n\n{text}"
+        prompt = (
+            f"<start_of_turn>user\n"
+            f"{user_content}<end_of_turn>\n"
+            f"<start_of_turn>model\n"
+        )
+
+        enable_mc_dropout(self.model)
+        pass_texts = []
+        for _ in range(n_passes):
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=2048,
+            ).to(self.device)
+
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=400,
+                    temperature=0.7,
+                    do_sample=True,
+                    top_p=0.9,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    eos_token_id=[self.tokenizer.eos_token_id, 107],
+                )
+            generated = self.tokenizer.decode(
+                outputs[0][inputs["input_ids"].shape[1]:],
+                skip_special_tokens=True,
+            ).strip()
+            pass_texts.append(generated)
+
+        disable_mc_dropout(self.model)
+        return aggregate_beliefs(pass_texts, n_passes)
+
+    def __call__(self, data):
+        inputs = data.get("inputs", "")
+        parameters = data.get("parameters", {})
+        mode = parameters.get("mode", "predict")
+        mc_passes = int(parameters.get("mc_passes", 0))
+
+        # Diagnostic: return handler info when requested
+        if parameters.get("info"):
+            return {
+                "handler_version": "2.0-mc",
+                "mc_dropout_available": self.mc_dropout_available,
+                "device": str(self.device),
+            }
+
+        if mode == "batch":
+            chunks = self._chunk_narrative(inputs)
+            results = []
+            for i, chunk in enumerate(chunks):
+                if mc_passes > 0 and self.mc_dropout_available:
+                    mc_beliefs = self._run_mc_inference(chunk, mc_passes)
+                    results.append({
+                        "chunk_index": i,
+                        "mc_beliefs": mc_beliefs,
+                        "mc_passes": mc_passes,
+                    })
+                else:
+                    generated = self._run_inference(chunk)
+                    results.append({"chunk_index": i, "generated_text": generated})
+            return {"total_chunks": len(chunks), "results": results}
+
+        if mc_passes > 0 and self.mc_dropout_available:
+            mc_beliefs = self._run_mc_inference(inputs, mc_passes)
+            return {"mc_beliefs": mc_beliefs, "mc_passes": mc_passes}
+
+        generated = self._run_inference(inputs)
+        return {"generated_text": generated}