refactor structure

app.py

@@ -1,61 +1,166 @@
 import os
 import re
+import gc
+from pathlib import Path
+
 import torch
 import gradio as gr
 
 from model_utils import load_model_and_tokenizer, generate_completion
-from retriever_stub import retrieve_code_stub
 
-# =========================
-# Model loading (HF Spaces compatible)
-# =========================
-if os.path.exists("reacc_generator/checkpoint-best"):
-    MODEL_PATH = "reacc_generator/checkpoint-best"
-else:
-    MODEL_PATH = "reacc_generator/checkpoint-last"
+
+# ============================================================
+# Path config
+# ============================================================
+
+BASE_DIR = Path(__file__).parent
+
+MODEL_PATHS = {
+    "Generator - Baseline": BASE_DIR / "checkpoint-best" / "baseline",
+    "Generator - EOL": BASE_DIR / "checkpoint-best" / "eol",
+}
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer, model = load_model_and_tokenizer(MODEL_PATH)
-model.to(device)
-model.eval()
 
-# =========================
+_current_model_name = None
+_current_tokenizer = None
+_current_model = None
+
+
+# ============================================================
+# Model loading
+# ============================================================
+
+def get_model(model_name: str):
+    """
+    Lazily load selected model.
+    Only one model is kept in memory at a time.
+    """
+
+    global _current_model_name, _current_tokenizer, _current_model
+
+    if model_name not in MODEL_PATHS:
+        raise ValueError(f"Unknown model option: {model_name}")
+
+    model_path = MODEL_PATHS[model_name]
+
+    if not model_path.exists():
+        raise FileNotFoundError(
+            f"Model path not found: {model_path}\n"
+            f"Expected structure: checkpoint-best/baseline and checkpoint-best/eol"
+        )
+
+    # Reuse current loaded model
+    if _current_model_name == model_name and _current_model is not None:
+        return _current_tokenizer, _current_model, model_path
+
+    # Unload old model if switching
+    if _current_model is not None:
+        del _current_model
+        del _current_tokenizer
+        _current_model = None
+        _current_tokenizer = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    print(f"Loading model: {model_name} from {model_path}")
+
+    tokenizer, model = load_model_and_tokenizer(str(model_path))
+    model.to(device)
+    model.eval()
+
+    _current_model_name = model_name
+    _current_tokenizer = tokenizer
+    _current_model = model
+
+    return tokenizer, model, model_path
+
+
+# ============================================================
 # Soft normalization adapters
-# =========================
+# ============================================================
+
+def normalize_line(line: str) -> str:
+    """
+    Soft-normalize one line to be closer to training token style.
+    Example:
+      def add(a, b):
+    becomes:
+      def add ( a , b ) :
+    """
+
+    # Put spaces around common Python punctuation/operators
+    line = re.sub(r"([()\[\]{}:,.=+\-*/<>])", r" \1 ", line)
+
+    # Collapse spaces
+    line = re.sub(r"\s+", " ", line)
+
+    return line.strip()
 
 
-def python_to_tokens(code: str) -> str:
+def context_to_tokens(code: str) -> str:
     """
-    Soft-normalize Python code into a format closer to training distribution
-    (spacing + <EOL>), without faking STR_LIT / NUM_LIT.
+    Convert normal-looking code into training-style token text.
+
+    Important:
+    - Preserve line boundaries as <EOL>
+    - Do not fake <STR_LIT> / <NUM_LIT>
     """
+
     code = code.replace("\t", "    ")
-    code = re.sub(r"([():,=+*/-])", r" \1 ", code)
-    code = re.sub(r"\s+", " ", code)
-    code = code.replace("\n", " <EOL> ")
-    return code.strip()
+    lines = code.splitlines()
+
+    normalized_lines = []
+    for line in lines:
+        norm = normalize_line(line)
+        if norm:
+            normalized_lines.append(norm)
+
+    return " <EOL> ".join(normalized_lines).strip()
 
 
-def tokens_to_python(code: str) -> str:
-    """Convert model token output back to readable Python."""
+def tokens_to_readable(code: str) -> str:
+    """
+    Convert generated token text back to readable form.
+    This is demo-level detokenization, not a perfect Python formatter.
+    """
+
     code = code.replace("<EOL>", "\n")
-    code = re.sub(r"\s*([():,=+*/-])\s*", r"\1", code)
+
+    # Remove spaces before punctuation
+    code = re.sub(r"\s+([)\]\}:,])", r"\1", code)
+
+    # Remove spaces after opening punctuation
+    code = re.sub(r"([(\[\{])\s+", r"\1", code)
+
+    # Compact common binary operators mildly
+    code = re.sub(r"\s*=\s*", " = ", code)
+    code = re.sub(r"\s*\+\s*", " + ", code)
+    code = re.sub(r"\s*-\s*", " - ", code)
+    code = re.sub(r"\s*\*\s*", " * ", code)
+    code = re.sub(r"\s*/\s*", " / ", code)
+    code = re.sub(r"\s*<\s*", " < ", code)
+    code = re.sub(r"\s*>\s*", " > ", code)
+
+    # Clean repeated spaces
+    code = re.sub(r"[ \t]+", " ", code)
+
     return code.strip()
 
-# =========================
+
+# ============================================================
 # Inference
-# =========================
+# ============================================================
 
+def run_demo(model_name: str, context: str):
+    tokenizer, model, model_path = get_model(model_name)
 
-def run_demo(context: str, use_retriever: bool):
-    # Raw -> normalized tokens
-    token_context = python_to_tokens(context)
+    token_context = context_to_tokens(context)
 
-    # Retriever (optional)
-    retrieved_raw = retrieve_code_stub(context) if use_retriever else ""
-    token_retrieved = python_to_tokens(retrieved_raw) if retrieved_raw else ""
+    # No retriever for now
+    token_retrieved = ""
 
-    # Generator
     token_output = generate_completion(
         model=model,
         tokenizer=tokenizer,
@@ -68,12 +173,12 @@ def run_demo(context: str, use_retriever: bool):
         stop_strings=["<EOL>"],
     )
 
-    # Tokens -> Python
-    python_output = tokens_to_python(token_output)
+    prediction = tokens_to_readable(token_output)
 
-    # Logs for explanation
     logs = (
-        "=== NORMALIZATION & GENERATION LOGS ===\n\n"
+        "=== DEMO LOGS ===\n\n"
+        f"[Selected model]\n{model_name}\n\n"
+        f"[Model path]\n{model_path}\n\n"
         "[Raw Context]\n"
         f"{context}\n\n"
         "[Context → Tokens]\n"
@@ -81,34 +186,40 @@ def run_demo(context: str, use_retriever: bool):
         "[Retrieved → Tokens]\n"
         f"{token_retrieved}\n\n"
         "[Generator Output → Tokens]\n"
-        f"{token_output}\n"
+        f"{token_output}\n\n"
+        "[Prediction]\n"
+        f"{prediction}\n"
     )
 
-    mode = "ReACC (Retriever + Generator)" if use_retriever else "Generator-only baseline"
-    return python_output, retrieved_raw, logs, mode
+    return prediction, logs
 
 
-# =========================
+# ============================================================
 # Gradio UI
-# =========================
+# ============================================================
+
 demo = gr.Interface(
     fn=run_demo,
     inputs=[
-        gr.Textbox(lines=12, label="Context (unfinished code)",
-                   placeholder="def sum(a, b):\n    "),
-        gr.Checkbox(label="Use Retriever (ReACC mode)", value=False),
+        gr.Dropdown(
+            choices=["Generator - Baseline", "Generator - EOL"],
+            value="Generator - Baseline",
+            label="Model",
+        ),
+        gr.Textbox(
+            lines=12,
+            label="Context",
+            placeholder="def add(a, b):\n    return",
+        ),
     ],
     outputs=[
         gr.Textbox(lines=8, label="Prediction"),
-        gr.Textbox(lines=6, label="Retrieved code"),
-        gr.Textbox(lines=12, label="Logs"),
-        gr.Textbox(lines=1, label="Mode"),
+        gr.Textbox(lines=14, label="Logs"),
     ],
-    title="ReACC Code Completion Demo",
+    title="ReACC Generator Demo",
     description=(
-        "User-friendly demo: input normal Python code.\n"
-        "The system softly normalizes input to the training-style tokens.\n"
-        "Logs show internal transformations for explanation."
+        "Compare Generator baseline and Generator + EOL. "
+        "Retriever integration will be added later."
     ),
 )
 

checkpoint-best/eol/config.json

@@ -0,0 +1,37 @@
+{
+  "_num_labels": 2,
+  "activation_function": "gelu_new",
+  "add_cross_attention": false,
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 0,
+  "dtype": "float32",
+  "embd_pdrop": 0.1,
+  "eos_token_id": 2,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": null,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "output_past": true,
+  "pad_token_id": 1,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.0.0",
+  "use_cache": true,
+  "vocab_size": 50007
+}

checkpoint-best/eol/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "pad_token_id": 1,
+  "transformers_version": "5.0.0"
+}

checkpoint-best/eol/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a300ae726608e2d94265138f1ccce63e08122e979e90cdb728ec798f19f54c38
+size 497006208

checkpoint-best/eol/tokenizer_config.json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<RET>",
+    "</RET>",
+    "<CTX>",
+    "</CTX>",
+    "<GEN>"
+  ],
+  "full_tokenizer_file": null,
+  "is_local": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sep_token": "<EOL>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|UNKNOWN|>"
+}

reacc_generator/checkpoint-best/.gitkeep

@@ -1 +0,0 @@
-.gitkeep