replicate railway

app.py

@@ -1,8 +1,9 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import json, math, torch, numpy as np
 from collections import Counter
+from typing import Optional
 import torch.nn as nn
 from huggingface_hub import hf_hub_download
 
@@ -16,19 +17,22 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ── Load assets ───────────────────────────────────────────────────────────────
-print("Loading model...")
-config   = json.load(open(hf_hub_download(HF_REPO_ID, "config.json")))
+# ── Load word lists & configs ─────────────────────────────────────────────────
+print("Loading configs and word lists...")
+config    = json.load(open(hf_hub_download(HF_REPO_ID, "config.json")))
+rl_config = json.load(open(hf_hub_download(HF_REPO_ID, "rl_config.json")))
 ANSWERS  = json.load(open(hf_hub_download(HF_REPO_ID, "answers.json")))
 ALLOWED  = json.load(open(hf_hub_download(HF_REPO_ID, "allowed.json")))
+WORD2IDX = {w: i for i, w in enumerate(ALLOWED)}
 LETTERS  = "abcdefghijklmnopqrstuvwxyz"
 L2I      = {c: i for i, c in enumerate(LETTERS)}
-INPUT_DIM   = config["input_dim"]
-OUTPUT_DIM  = config["output_dim"]
-OPENING     = config["opening_guess"]
+INPUT_DIM  = config["input_dim"]
+OUTPUT_DIM = config["output_dim"]
+OPENING    = config["opening_guess"]
 WIN_PATTERN = (2, 2, 2, 2, 2)
 
-# ── Model ─────────────────────────────────────────────────────────────────────
+
+# ── Model architecture ────────────────────────────────────────────────────────
 class WordleNet(nn.Module):
     def __init__(self):
         super().__init__()
@@ -41,17 +45,50 @@ class WordleNet(nn.Module):
         )
     def forward(self, x): return self.net(x)
 
+
+class RLWordleNet(nn.Module):
+    """Same encoder as WordleNet but with BatchNorm-safe single-sample forward."""
+    def __init__(self):
+        super().__init__()
+        h = rl_config["hidden"]
+        self.encoder = nn.Sequential(
+            nn.Linear(INPUT_DIM, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(0.3),
+            nn.Linear(h, h),         nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(0.3),
+            nn.Linear(h, 256),       nn.BatchNorm1d(256), nn.ReLU(),
+        )
+        self.policy_head = nn.Linear(256, OUTPUT_DIM)
+
+    def forward(self, x):
+        single = x.shape[0] == 1
+        if single:
+            x = x.repeat(2, 1)
+        feat = self.encoder(x)
+        if single:
+            feat = feat[:1]
+        return self.policy_head(feat)
+
+
+# ── Load weights ──────────────────────────────────────────────────────────────
+print("Loading supervised model...")
 model = WordleNet()
 model.load_state_dict(
     torch.load(hf_hub_download(HF_REPO_ID, "model_weights.pt"), map_location="cpu")
 )
 model.eval()
-print("Model loaded ✅")
 
-# ── Helpers ───────────────────────────────────────────────────────────────────
+print("Loading RL model...")
+rl_model = RLWordleNet()
+rl_model.load_state_dict(
+    torch.load(hf_hub_download(HF_REPO_ID, "rl_model_weights.pt"), map_location="cpu"), strict=False
+)
+rl_model.eval()
+print("Both models loaded.")
+
+
+# ── Core logic ────────────────────────────────────────────────────────────────
 def get_pattern(guess, answer):
-    pattern = [0] * 5
-    counts  = Counter(answer)
+    pattern = [0]*5
+    counts = Counter(answer)
     for i in range(5):
         if guess[i] == answer[i]:
             pattern[i] = 2
@@ -68,105 +105,82 @@ def filter_words(words, guess, pattern):
 def entropy_score(guess, possible):
     buckets = Counter(get_pattern(guess, w) for w in possible)
     n = len(possible)
-    return sum(-(c / n) * math.log2(c / n) for c in buckets.values())
+    return sum(-(c/n)*math.log2(c/n) for c in buckets.values())
 
 def encode_board(history):
     vec = np.zeros(INPUT_DIM, dtype=np.float32)
     for word, pattern in history:
         for pos, (letter, state) in enumerate(zip(word, pattern)):
-            vec[L2I[letter] * 15 + pos * 3 + state] = 1.0
+            vec[L2I[letter]*15 + pos*3 + state] = 1.0
     return vec
 
-def is_consistent(word, history):
-    for guess, pattern in history:
-        green_letters = {letter for letter, state in zip(guess, pattern) if state == 2}
-        for pos, (letter, state) in enumerate(zip(guess, pattern)):
-            if state == 2:
-                if word[pos] != letter:
-                    return False
-            elif state == 1:
-                if letter not in word or word[pos] == letter:
-                    return False
-            else:
-                if letter not in green_letters and letter in word:
-                    return False
-    return True
-
-def model_suggest(history, possible):
-    if not possible:       return None
-    if len(possible) == 1: return possible[0]
-    if not history:        return OPENING
-
-    already_guessed = {w for w, _ in history}
-    possible_not_guessed = [w for w in possible if w not in already_guessed]
-
-    if len(possible) <= 6:
-        ambiguous = set()
-        for pos in range(5):
-            letters_at_pos = {w[pos] for w in possible}
-            if len(letters_at_pos) > 1:
-                ambiguous.update(letters_at_pos)
-
-        best_word, best_score = None, -1
-        for g in ALLOWED:
-            if g in already_guessed:
-                continue
-            if not is_consistent(g, history):
-                continue
-            if g in possible and len(possible) > 2:
-                continue
-            score = len(set(g) & ambiguous) * 2 + entropy_score(g, possible)
-            if score > best_score:
-                best_score, best_word = score, g
-
-        if not best_word:
-            best_word = possible_not_guessed[0] if possible_not_guessed else possible[0]
-        return best_word
+def get_logits(history, possible, use_rl=False):
+    """Get top-20 model words using constraint-filtered mask."""
+    active_model = rl_model if use_rl else model
+    possible_set = set(possible)
 
     state = torch.tensor(encode_board(history)).unsqueeze(0)
     with torch.no_grad():
-        logits = model(state)[0]
+        logits = active_model(state)[0]
 
-    top50 = [ALLOWED[i] for i in logits.topk(50).indices.tolist()]
-    valid = [w for w in top50
-             if w not in already_guessed and is_consistent(w, history)]
+    if use_rl:
+        mask = torch.full((OUTPUT_DIM,), float('-inf'))
+        for i, w in enumerate(ALLOWED):
+            if w in possible_set:
+                mask[i] = 0.0
+        if mask.max() == float('-inf'):
+            mask[:] = 0.0
+        logits = logits + mask
 
-    if not valid:
-        return max(possible_not_guessed or possible,
-                   key=lambda w: entropy_score(w, possible))
+    return [ALLOWED[i] for i in logits.topk(20).indices.tolist()]
 
-    return max(valid[:10], key=lambda w: entropy_score(w, possible))
 
+def model_suggest(history, possible, use_rl=False):
+    if not possible:       return None
+    if len(possible) == 1: return possible[0]
+    if not history:        return OPENING
+    guessed = {w for w, _ in history}
 
-def top_suggestions(history, possible, n=5):
-    if not possible: return []
+    model_words = get_logits(history, possible, use_rl)
+
+    if len(possible) <= 6:
+        best, best_worst = None, float('inf')
+        for g in list(possible) + model_words:
+            if g in guessed: continue
+            worst = max(Counter(get_pattern(g, w) for w in possible).values())
+            if worst < best_worst:
+                best_worst, best = worst, g
+        return best or possible[0]
 
-    already_guessed = {w for w, _ in history}
+    candidates = list(dict.fromkeys(model_words + list(possible)))
+    candidates = [w for w in candidates if w not in guessed]
+    if not candidates: return possible[0]
+    return max(candidates, key=lambda w: entropy_score(w, possible))
 
+
+def top_suggestions(history, possible, use_rl=False, n=5):
+    if not possible: return []
+    guessed = {w for w, _ in history}
     if not history:
-        candidates = [OPENING] + [w for w in ALLOWED if w != OPENING][:30]
+        candidates = [OPENING] + [w for w in ALLOWED if w != OPENING][:20]
     else:
-        state = torch.tensor(encode_board(history)).unsqueeze(0)
-        with torch.no_grad():
-            logits = model(state)[0]
-        candidates = [ALLOWED[i] for i in logits.topk(50).indices.tolist()]
-
-    candidates = [w for w in candidates
-                  if w not in already_guessed and is_consistent(w, history)]
+        model_words = get_logits(history, possible, use_rl)
+        candidates  = list(dict.fromkeys(model_words + list(possible)))
 
     possible_set = set(possible)
-    scored = [
-        {
-            "word":        w,
-            "entropy":     round(entropy_score(w, possible), 3),
-            "is_possible": w in possible_set,
-        }
-        for w in candidates
-    ]
-    scored.sort(key=lambda x: (-x["entropy"], not x["is_possible"]))
+    candidates = [w for w in candidates if w in possible_set and w not in guessed]
+
+    # fallback — if all possible words were guessed, show from full possible
+    if not candidates:
+        candidates = [w for w in possible if w not in guessed]
+
+    scored = [{"word": w, "entropy": round(entropy_score(w, possible), 3), "is_possible": True}
+              for w in candidates]
+    scored.sort(key=lambda x: -x["entropy"])
     return scored[:n]
 
-# ── Models ────────────────────────────────────────────────────────────────────
+
+# ── Schemas ───────────────────────────────────────────────────────────────────
 class GuessEntry(BaseModel):
     word: str
     pattern: list[int]
@@ -181,14 +195,20 @@ class SuggestResponse(BaseModel):
     bits_remaining: float
     solved: bool
     message: str
+    model_used: str
 
-# ── Routes ────────────────────────────────────────────────────────────────────
+
+# ── Endpoints ─────────────────────────────────────────────────────────────────
 @app.get("/")
 def root():
     return {"status": "ok", "opener": OPENING}
 
 @app.post("/suggest", response_model=SuggestResponse)
-def suggest(req: SuggestRequest):
+def suggest(
+    req: SuggestRequest,
+    model: str = Query(default="supervised", pattern="^(supervised|rl)$")
+):
+    use_rl   = model == "rl"
     possible = list(ANSWERS)
 
     for entry in req.history:
@@ -196,12 +216,12 @@ def suggest(req: SuggestRequest):
         pattern = tuple(entry.pattern)
         if len(word) != 5:
             raise HTTPException(400, f"Word must be 5 letters: {word}")
-        if len(pattern) != 5 or not all(p in (0, 1, 2) for p in pattern):
+        if len(pattern) != 5 or not all(p in (0,1,2) for p in pattern):
             raise HTTPException(400, "Pattern must be 5 values of 0, 1, or 2")
         if pattern == WIN_PATTERN:
             return SuggestResponse(
                 suggestion=word, top_suggestions=[], possible_count=1,
-                bits_remaining=0.0, solved=True,
+                bits_remaining=0.0, solved=True, model_used=model,
                 message=f"Solved in {len(req.history)} guesses!"
             )
         possible = filter_words(possible, word, pattern)
@@ -210,10 +230,10 @@ def suggest(req: SuggestRequest):
         raise HTTPException(422, "No possible words remaining. Check your pattern input.")
 
     history_tuples = [(e.word.lower(), tuple(e.pattern)) for e in req.history]
-    suggestion     = model_suggest(history_tuples, possible)
+    suggestion     = model_suggest(history_tuples, possible, use_rl=use_rl)
     if not suggestion:
         suggestion = possible[0]
-    top_suggs      = top_suggestions(history_tuples, possible)
+    top_suggs      = top_suggestions(history_tuples, possible, use_rl=use_rl)
     bits           = math.log2(len(possible)) if len(possible) > 1 else 0.0
 
     return SuggestResponse(
@@ -222,9 +242,14 @@ def suggest(req: SuggestRequest):
         possible_count=len(possible),
         bits_remaining=round(bits, 2),
         solved=False,
+        model_used=model,
         message=f"{len(possible)} words remaining — try {suggestion.upper()}"
     )
 
 @app.get("/opener")
 def get_opener():
     return {"word": OPENING}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)