Update app.py

app.py

@@ -1,90 +1,102 @@
-from retriever_stub import retrieve_code_stub
-from model_utils import load_model_and_tokenizer, generate_completion
-import sys
-from pathlib import Path
+import os
 import torch
 import gradio as gr
 
-# ===== Path setup =====
-APP_DIR = Path(__file__).parent
-GENERATOR_DIR = APP_DIR.parent
-CHECKPOINT_DIR = GENERATOR_DIR / "reacc_generator"
-
-sys.path.insert(0, str(GENERATOR_DIR))
-
-
-# ===== Load model once (important for demo speed) =====
-# MODEL_PATH = CHECKPOINT_DIR / "checkpoint-best"
-# if not MODEL_PATH.exists():
-#     MODEL_PATH = CHECKPOINT_DIR / "checkpoint-last"
+from model_utils import load_model_and_tokenizer, generate_completion
+from retriever_stub import retrieve_code_stub
 
-import os
 
+# =========================
+# Model loading
+# =========================
 if os.path.exists("reacc_generator/checkpoint-best"):
     MODEL_PATH = "reacc_generator/checkpoint-best"
 else:
     MODEL_PATH = "reacc_generator/checkpoint-last"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer, model = load_model_and_tokenizer(str(MODEL_PATH))
+tokenizer, model = load_model_and_tokenizer(MODEL_PATH)
 model.to(device)
 model.eval()
 
-print("Model loaded from:", MODEL_PATH)
-print("Device:", device)
+
+# =========================
+# Demo-level token adapter
+# =========================
+def python_to_demo_tokens(code: str) -> str:
+    """Convert Python code to a simplified token format for the model."""
+    return code.replace("\n", " <EOL> ")
 
 
-# ===== Core inference logic =====
-def run_demo(context: str, use_retriever: bool):
-    """
-    context: unfinished code (tokenized style)
-    use_retriever: toggle ON/OFF
-    """
+def demo_tokens_to_python(code: str) -> str:
+    """Convert model token output back to readable Python code."""
+    return code.replace("<EOL>", "\n")
 
-    retrieved = ""
 
-    if use_retriever:
-        retrieved = retrieve_code_stub(context)
+# =========================
+# Inference function
+# =========================
+def run_demo(python_context: str, use_retriever: bool):
 
-    prediction = generate_completion(
+    # 1. Python → tokenized (hidden from user)
+    token_context = python_to_demo_tokens(python_context)
+
+    # 2. Retriever (optional)
+    retrieved = retrieve_code_stub(python_context) if use_retriever else ""
+    token_retrieved = python_to_demo_tokens(retrieved) if retrieved else ""
+
+    # 3. Generator
+    token_output = generate_completion(
         model=model,
         tokenizer=tokenizer,
-        retrieved=retrieved,
-        context=context,
+        retrieved=token_retrieved,
+        context=token_context,
         device=device,
         max_length=256,
-        max_new_tokens=16,   # keep short for demo
+        max_new_tokens=16,
         do_sample=False,
-        stop_strings=["<EOL>"]
+        stop_strings=["<EOL>"],
     )
 
-    mode = "Retriever + Generator (ReACC)" if use_retriever else "Generator-only baseline"
+    # 4. Token → Python
+    python_output = demo_tokens_to_python(token_output)
+
+    # 5. Logs
+    logs = (
+        "=== TOKENIZATION LOGS ===\n\n"
+        "[Input → Tokens]\n"
+        f"{token_context}\n\n"
+        "[Retrieved → Tokens]\n"
+        f"{token_retrieved}\n\n"
+        "[Generator Output → Tokens]\n"
+        f"{token_output}\n"
+    )
 
-    return prediction, retrieved, mode
+    mode = "ReACC (Retriever + Generator)" if use_retriever else "Generator-only baseline"
+    return python_output.strip(), retrieved, logs, mode
 
 
-# ===== Gradio UI =====
+# =========================
+# Gradio UI
+# =========================
 demo = gr.Interface(
     fn=run_demo,
     inputs=[
-        gr.Textbox(
-            lines=12,
-            label="Context (unfinished code)",
-            placeholder="Paste code context here (tokenized style: <EOL>, <STR_LIT>, <NUM_LIT>)"
-        ),
+        gr.Textbox(lines=12, label="Python context (unfinished code)",
+                   placeholder="def sum(a, b):\n    "),
         gr.Checkbox(label="Use Retriever (ReACC mode)", value=False),
     ],
     outputs=[
-        gr.Textbox(lines=6, label="Prediction"),
+        gr.Textbox(lines=8, label="Prediction (Python code)"),
         gr.Textbox(lines=6, label="Retrieved code"),
+        gr.Textbox(lines=12, label="Logs (tokenization & generation)"),
         gr.Textbox(lines=1, label="Mode"),
     ],
     title="ReACC Code Completion Demo",
     description=(
-        "Toggle Retriever ON/OFF to compare:\n"
-        "- Generator-only baseline\n"
-        "- Retriever-augmented generation (ReACC)\n\n"
-        "This demo runs on a fine-tuned CodeGPT generator."
+        "Enter normal Python code.\n"
+        "The system will internally tokenize it for the generator.\n"
+        "You can view tokenization and generation details in the Logs section."
     ),
 )