add json-repair

.gitignore

@@ -2,4 +2,5 @@
 local_test.py
 .env
 upload_to_r2.py
-Build
+Build
+repair_test.py

app.py

@@ -18,9 +18,12 @@ import boto3
 from typing import Generator
 
 from prompts import get_system_prompt
-from helpers import parse_dialogue_tags, create_episode_zip, postprocess_script
+from helpers import parse_dialogue_tags, create_episode_zip, postprocess_script, expected_schema
 from forced_alignment import forced_align
 
+import json_repair
+import re
+
 llm_model_path = hf_hub_download(
     repo_id="unsloth/gemma-4-26B-A4B-it-GGUF",
     filename="gemma-4-26B-A4B-it-UD-Q4_K_M.gguf"
@@ -68,6 +71,7 @@ def warmup():
 
 
 
+
 @spaces.GPU()
 def run_episode_generation_pipeline(character_names: dict[str, str], problem: str, extras: str = "", language="English"):
     llm = Llama(model_path=llm_model_path, n_gpu_layers=-1, n_ctx=8192, flash_attn=True, verbose=False)
@@ -93,7 +97,17 @@ def run_episode_generation_pipeline(character_names: dict[str, str], problem: st
 
     # Episode JSON
     try:
-        script = json.loads(script_raw)
+        # Try fixing the script if broken
+        script = json_repair.loads(
+        script_raw,
+        schema=expected_schema,
+        schema_repair_mode="salvage"
+        )
+        # Fix starting look_at tags
+        for item in script.get("dialogues", []):
+            original_text = item.get("dialogue", "")
+            fixed_text = re.sub(r'(?<!<)look_at:', '<look_at:', original_text)
+            item["dialogue"] = fixed_text
     except json.JSONDecodeError as e:
         print(f"JSON Parse Error: {e}")
         yield "Error: LLM generated an invalid script format. Please try again."

helpers.py

@@ -64,4 +64,25 @@ def postprocess_script(script, character_names: dict[str, str]):
             if target_name in name_to_role:
                 look_target["character"] = name_to_role[target_name]
         
-    return script
+    return script
+
+
+
+expected_schema = {
+    "type": "object",
+    "properties": {
+        "dialogues": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "character": {"type": "string"},
+                    "dialogue": {"type": "string"},
+                    "expression": {"type": "string"}
+                },
+                "required": ["character", "dialogue", "expression"]
+            }
+        }
+    },
+    "required": ["dialogues"]
+}

requirements.txt

@@ -12,4 +12,6 @@ hf_transfer
 omnivoice
 --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu124
 llama-cpp-python
-boto3
+boto3
+json-repair
+jsonschema