Upload infer.py

infer.py

@@ -10,12 +10,7 @@ SPACE_IDX = CHAR_TO_INT.get(' ', 0)
 CONTEXT_SIZE = 64
 MAX_GEN = 60
 
-def generate(model, prompt_text):
-    clean = prompt_text.replace(' ', '')
-    stripped = ''.join(c for c in clean if c in CHAR_TO_INT)
-    if '=' not in stripped:
-        return prompt_text + '?'
-    prompt = stripped[:stripped.index('=') + 1]
+def _run(model, prompt):
     generated = ''
     for _ in range(MAX_GEN):
         prefix = (prompt + generated).rjust(CONTEXT_SIZE, ' ')[:CONTEXT_SIZE]
@@ -28,7 +23,48 @@ def generate(model, prompt_text):
         if not char or char == ' ' or char == '|':
             break
         generated += char
-    return prompt_text[:prompt_text.index('=') + 1] + generated
+    return prompt + generated
+
+def _count_operators(expr):
+    count = 0
+    for c in expr:
+        if c in '+-*':
+            count += 1
+    return count
+
+def _strip_outer_parens(expr):
+    while expr and expr[0] == '(' and expr[-1] == ')':
+        depth = 0
+        balanced = True
+        for i, c in enumerate(expr):
+            if c == '(':
+                depth += 1
+            elif c == ')':
+                depth -= 1
+            if depth == 0 and i < len(expr) - 1:
+                balanced = False
+                break
+        if balanced:
+            expr = expr[1:-1]
+        else:
+            break
+    return expr
+
+def generate(model, prompt_text):
+    clean = prompt_text.replace(' ', '')
+    stripped = ''.join(c for c in clean if c in CHAR_TO_INT)
+    if '=' not in stripped:
+        return prompt_text + '?'
+    prompt = stripped[:stripped.index('=') + 1]
+    expr = _strip_outer_parens(stripped[:stripped.index('=')])
+    if _count_operators(expr) == 1:
+        wrapped = f"0+({expr})="
+        chain = _run(model, wrapped)
+        parts = chain.split('=')
+        answer = parts[-1].rstrip('| ')
+        if answer and answer != expr and not answer.startswith('-'):
+            return prompt_text[:prompt_text.index('=') + 1] + answer
+    return _run(model, prompt)
 
 def main():
     if len(sys.argv) < 2: