Update app.py

app.py

@@ -136,45 +136,38 @@ class Transformer(nn.Module):
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load tokenizers
-sp_pseudo = spm.SentencePieceProcessor(model_file="pseudocode_tokenizer.model")  # For decoding pseudocode (target)
-sp_code = spm.SentencePieceProcessor(model_file="code_tokenizer.model")      # For encoding C++ (source)
+sp_pseudo = spm.SentencePieceProcessor(model_file="pseudo.model")  # For decoding pseudocode (target)
+sp_code = spm.SentencePieceProcessor(model_file="code.model")      # For encoding C++ (source)
 
 # Load the full saved model (architecture + weights)
-model_path = "code2pseudo.pth"
+model_path = "transformer_cpp_to_pseudo.pth"
 model = torch.load(model_path, map_location=device, weights_only=False)
 model.eval()
 model = model.to(device)
 
-def generate_pseudocode(cpp_code, max_len=500):
+def generate_pseudocode(cpp_code, max_len):
     """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>
-    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
-    
+    response = ""
     with torch.no_grad():
-        for i in range(max_len):
+        for _ 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 == eos_id:  # Stop at <EOS>
-                print("EOS detected, stopping generation.")
+            if next_token == 3:  # <END>=3 (adjust if your EOS ID differs)
                 break
-        print("Generation complete or max length reached.")
-    
     yield response  # Final output
 
-def generate_output(cpp_code):
-    """Wrapper for Gradio interface with streaming."""
-    for response in generate_pseudocode(cpp_code, max_len=500):
+def respond(message, history, max_tokens):
+    """Wrapper for Gradio interface."""
+    for response in generate_pseudocode(message, max_tokens):
         yield response
 
 # Gradio UI setup with Blocks
@@ -186,28 +179,17 @@ with gr.Blocks(title="C++ to Pseudocode Transformer") as demo:
         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")
+    submit_btn = gr.Button("Submit", variant="primary")
     pseudocode_output = gr.Textbox(
         label="Generated Pseudocode",
         lines=5
     )
     
     submit_btn.click(
-        fn=generate_output,
-        inputs=[cpp_input],
+        fn=respond,
+        inputs=[cpp_input, gr.State(value=[]), gr.Slider(minimum=10, maximum=1000, value=50, step=1, visible=False)],
         outputs=pseudocode_output
     )
 
-demo.launch()
-
-# Custom CSS
-demo.css = """
-.btn-blue {
-    background-color: #007bff;
-    color: white;
-    border: none;
-}
-.btn-blue:hover {
-    background-color: #0056b3;
-}
-"""
+if __name__ == "__main__":
+    demo.launch()