Update app.py

app.py

@@ -13,6 +13,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # CONFIG
 # =========================================================
 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+MAX_NEW_TOKENS = 256
 
 app = FastAPI(title="Japanese Corrector API")
 
@@ -34,23 +35,19 @@ class InferenceRequest(BaseModel):
     input: str
 
 # =========================================================
-# SYSTEM PROMPT
+# SYSTEM PROMPT (STRICT)
 # =========================================================
 SYSTEM_PROMPT = """
 You are a professional Japanese language teacher.
-Input:
-- Japanese written in ROMAJI.
-Tasks:
-1. Correct the sentence into natural Japanese.
-2. Use appropriate Kanji.
-3. Provide full Hiragana reading.
-4. Explain grammar changes briefly.
-Rules:
-- Do NOT change meaning or polarity.
-- Do NOT invent verb forms.
-- Output ONLY valid JSON object (single JSON object) and nothing else.
-- No markdown, no extra text.
-Output format:
+
+Return ONLY valid JSON.
+No text before or after JSON.
+No markdown.
+No explanations outside JSON.
+
+The JSON MUST be fully closed.
+
+Schema:
 {
   "kanji": "...",
   "hiragana": "...",
@@ -59,16 +56,13 @@ Output format:
 """.strip()
 
 # =========================================================
-# JSON EXTRACTION & SAFE PARSING
+# JSON UTILITIES
 # =========================================================
 def _strip_code_fences(text: str) -> str:
-    # remove triple backticks and single backticks
     text = re.sub(r"```(?:json)?", "", text, flags=re.IGNORECASE)
-    text = text.replace("`", "")
-    return text
+    return text.replace("`", "")
 
 def _normalize_quotes(text: str) -> str:
-    # convert smart quotes to ASCII quotes
     replacements = {
         "“": '"', "”": '"', "‘": "'", "’": "'",
         "（": "(", "）": ")", "：": ":", "，": ","
@@ -78,73 +72,68 @@ def _normalize_quotes(text: str) -> str:
     return text
 
 def _remove_trailing_commas(text: str) -> str:
-    # remove trailing commas before } or ]
-    text = re.sub(r",\s*(?=[}\]])", "", text)
-    return text
+    return re.sub(r",\s*(?=[}\]])", "", text)
 
 def _find_balanced_json(text: str) -> Optional[str]:
-    """Find first balanced JSON object substring using brace counting."""
-    text = _strip_code_fences(text)
-    text = _normalize_quotes(text)
-    # find first '{'
+    text = _normalize_quotes(_strip_code_fences(text))
     for m in re.finditer(r"\{", text):
         start = m.start()
         depth = 0
         for i in range(start, len(text)):
-            ch = text[i]
-            if ch == "{":
+            if text[i] == "{":
                 depth += 1
-            elif ch == "}":
+            elif text[i] == "}":
                 depth -= 1
                 if depth == 0:
-                    candidate = text[start:i + 1]
-                    return candidate
+                    return text[start:i + 1]
     return None
 
 def safe_parse_json(text: str) -> Dict:
-    """
-    Robustly extract and parse the first JSON object in `text`.
-    Attempts:
-      1) Balanced-brace extraction + json.loads
-      2) Apply small repairs (remove trailing commas, normalize quotes) + json.loads
-      3) ast.literal_eval on candidate (accepts single quotes / python-style)
-    Raises ValueError if parsing fails.
-    """
-    raw = text or ""
-    candidate = _find_balanced_json(raw)
+    candidate = _find_balanced_json(text)
     if not candidate:
-        raise ValueError("No JSON-looking substring found")
+        raise ValueError("No JSON found")
 
-    # try direct json.loads
     try:
         return json.loads(candidate)
-    except Exception as e1:
-        # try small repair: remove trailing commas
+    except Exception:
         repaired = _remove_trailing_commas(candidate)
         try:
             return json.loads(repaired)
-        except Exception as e2:
-            # try to fallback to ast.literal_eval (accepts single quotes)
-            try:
-                obj = ast.literal_eval(repaired)
-                if isinstance(obj, dict):
-                    # ensure converted to json-compatible (strings) using dumps/loads
-                    return json.loads(json.dumps(obj, ensure_ascii=False))
-            except Exception as e3:
-                logger.debug("JSON parse attempts failed: %s | %s | %s", e1, e2, e3)
-                raise ValueError("JSON parsing failed") from e3
+        except Exception:
+            obj = ast.literal_eval(repaired)
+            if isinstance(obj, dict):
+                return json.loads(json.dumps(obj, ensure_ascii=False))
+            raise ValueError("JSON parsing failed")
+
+# =========================================================
+# GENERATION (WITH RETRY)
+# =========================================================
+def generate_once(prompt: str) -> str:
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=False,
+            temperature=0.0,
+            repetition_penalty=1.05,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    gen = output[0][inputs.input_ids.shape[1]:]
+    return tokenizer.decode(gen, skip_special_tokens=True).strip()
 
 # =========================================================
 # API ENDPOINT
 # =========================================================
 @app.post("/infer")
 def infer(req: InferenceRequest):
+
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": req.input.strip()}
     ]
 
-    # build the underlying prompt using tokenizer helper (keeps chat formatting)
     try:
         prompt = tokenizer.apply_chat_template(
             messages,
@@ -152,50 +141,40 @@ def infer(req: InferenceRequest):
             add_generation_prompt=True
         )
     except Exception:
-        # fallback: simple concatenation if apply_chat_template isn't available
         prompt = SYSTEM_PROMPT + "\n\nUser: " + req.input.strip() + "\nAssistant:"
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=False,
-            temperature=0.0,
-            repetition_penalty=1.05,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.eos_token_id,
-            # num_return_sequences=1 is default
-        )
-
-    generated = outputs[0][inputs.input_ids.shape[1]:]
-    raw_text = tokenizer.decode(generated, skip_special_tokens=True).strip()
-    logger.info("Raw model output: %s", raw_text[:1000])
+    # 🔁 TRY TWICE (important)
+    last_raw = ""
+    for attempt in range(2):
+        raw = generate_once(prompt)
+        last_raw = raw
+        logger.info("Raw model output (attempt %d): %s", attempt + 1, raw[:500])
 
-    try:
-        data = safe_parse_json(raw_text)
-        return {
-            "result": {
-                "kanji": data.get("kanji", ""),
-                "hiragana": data.get("hiragana", ""),
-                "explanation": data.get("explanation", ""),
-                "warnings": []
-            }
-        }
-    except Exception as exc:
-        # include raw_text as a debugging warning (helpful for diagnosing why parse failed)
-        return {
-            "result": {
-                "kanji": "",
-                "hiragana": "",
-                "explanation": "モデルの出力を解析できませんでした。",
-                "warnings": [
-                    "Unparseable model output",
-                    raw_text[:2000]  # limit length
-                ]
+        try:
+            data = safe_parse_json(raw)
+            return {
+                "result": {
+                    "kanji": data.get("kanji", ""),
+                    "hiragana": data.get("hiragana", ""),
+                    "explanation": data.get("explanation", ""),
+                    "warnings": []
+                }
             }
+        except Exception:
+            continue
+
+    # 🚨 Guaranteed safe fallback
+    return {
+        "result": {
+            "kanji": "",
+            "hiragana": "",
+            "explanation": "モデルの出力を解析できませんでした。",
+            "warnings": [
+                "Unparseable model output",
+                last_raw[:1000]
+            ]
         }
+    }
 
 # =========================================================
 # LOCAL RUN