Update app.py

app.py

@@ -145,42 +145,69 @@ model = torch.load(model_path, map_location=device, weights_only=False)
 model.eval()
 model = model.to(device)
 
-def generate_pseudocode(cpp_code, max_len):
+def generate_pseudocode(cpp_code, max_len=500):
     """Generate pseudocode from C++ code with streaming output."""
     model.eval()
     src = torch.tensor([sp_code.encode_as_ids(cpp_code)], dtype=torch.long, device=device)  # Tokenize C++ code
     tgt = torch.tensor([[2]], dtype=torch.long, device=device)  # <bos_id>=2
     
     generated_tokens = [2]  # Start with <START>
-    response = ""
+    eos_id = sp_pseudo.eos_id()  # Dynamically get <EOS> ID
+    print(f"Input C++ tokens: {sp_code.encode_as_ids(cpp_code)}")  # Debug input
+    print(f"Using EOS ID: {eos_id}")  # Debug EOS ID
+    
     with torch.no_grad():
-        for _ in range(max_len):
+        for i in range(max_len):
             output = model(src, tgt)
             next_token = output[:, -1, :].argmax(-1).item()
             generated_tokens.append(next_token)
             tgt = torch.cat([tgt, torch.tensor([[next_token]], device=device)], dim=1)
             response = sp_pseudo.decode_ids(generated_tokens)  # Decode to pseudocode
+            print(f"Step {i}: Next token = {next_token}, Partial output = {response}")  # Debug step
             yield response  # Yield partial output
-            if next_token == 3:  # <END>=3 (adjust if your EOS ID differs)
+            if next_token == eos_id:  # Stop at <EOS>
+                print("EOS detected, stopping generation.")
                 break
+        print("Generation complete or max length reached.")
+    
     yield response  # Final output
 
-def respond(message, history, max_tokens):
-    """Wrapper for Gradio interface."""
-    for response in generate_pseudocode(message, max_tokens):
+def generate_output(cpp_code):
+    """Wrapper for Gradio interface with streaming."""
+    for response in generate_pseudocode(cpp_code, max_len=500):
         yield response
 
-# Gradio interface
-demo = gr.ChatInterface(
-    respond,
-    chatbot=gr.Chatbot(label="C++ to Pseudocode Generator"),
-    textbox=gr.Textbox(placeholder="Enter C++ code (e.g., 'int x = 5; for(int i=0; i<x; i++) cout << i;')", label="C++ Code"),
-    additional_inputs=[
-        gr.Slider(minimum=10, maximum=1000, value=50, step=1, label="Max tokens"),
-    ],
-    title="C++ to Pseudocode Transformer",
-    description="Convert C++ code to pseudocode using a custom transformer trained on the SPoC dataset.",
-)
-
-if __name__ == "__main__":
-    demo.launch()
+# Gradio UI setup with Blocks
+with gr.Blocks(title="C++ to Pseudocode Transformer") as demo:
+    gr.Markdown("## C++ to Pseudocode Converter")
+    gr.Markdown("Enter C++ code below and press Submit to generate pseudocode.")
+    cpp_input = gr.Textbox(
+        label="C++ Code",
+        placeholder="e.g., 'int x = 5; for(int i=0; i<x; i++) cout << i;'",
+        lines=5
+    )
+    submit_btn = gr.Button("Submit", variant="primary", elem_classes="btn-blue")
+    pseudocode_output = gr.Textbox(
+        label="Generated Pseudocode",
+        lines=5
+    )
+    
+    submit_btn.click(
+        fn=generate_output,
+        inputs=[cpp_input],
+        outputs=pseudocode_output
+    )
+
+demo.launch()
+
+# Custom CSS
+demo.css = """
+.btn-blue {
+    background-color: #007bff;
+    color: white;
+    border: none;
+}
+.btn-blue:hover {
+    background-color: #0056b3;
+}
+"""