Update app.py

app.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import sentencepiece as spm
 import math
 
-# Transformer class definitions (unchanged)
+# Define Transformer components (unchanged)
 class MultiHeadAttention(nn.Module):
     def __init__(self, d_model, num_heads):
         super(MultiHeadAttention, self).__init__()
@@ -132,82 +132,55 @@ class Transformer(nn.Module):
         output = self.fc(dec_output)
         return output
 
-# Device setup
+# Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load tokenizers
-sp_code = spm.SentencePieceProcessor(model_file="code_tokenizer.model")      # C++ tokenizer for input
-sp_pseudo = spm.SentencePieceProcessor(model_file="pseudocode_tokenizer.model")  # Pseudocode tokenizer for output
+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)
 
-# Load the model trained for C++ to pseudocode
-model_path = "c2p.pth"  # Ensure this is the correct model for C++ to pseudocode
+# Load the full saved model (architecture + weights)
+model_path = "c2p.pth"
 model = torch.load(model_path, map_location=device, weights_only=False)
 model.eval()
 model = model.to(device)
 
-# Function to generate pseudocode from C++ code with streaming
-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++ input
-    tgt = torch.tensor([[2]], dtype=torch.long, device=device)  # <BOS> token (ID=2)
-    
-    generated_tokens = [2]  # Start with <BOS>
-    eos_id = sp_pseudo.eos_id()  # Dynamically get <EOS> ID from tokenizer
-    print(f"Input C++ tokens: {sp_code.encode_as_ids(cpp_code)}")  # Debug input
-    print(f"Using EOS ID: {eos_id}")  # Debug EOS ID
+    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 = ""
     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 for streaming
-            if next_token == eos_id:  # Stop at <EOS>
-                print("EOS detected, stopping generation.")
+            yield response  # Yield partial output
+            if next_token == 3:  # <END>=3 (adjust if your EOS ID differs)
                 break
-        print("Generation complete or max length reached.")
-    
     yield response  # Final output
 
-# Gradio interface function with streaming
-def generate_output(cpp_code):
-    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 gr.Blocks(title="C++ to Pseudocode Transformer") as demo:
-    gr.Markdown("## C++ to Pseudocode Converter")
-    gr.Markdown("Enter C++ code below to generate pseudocode.")
-    cpp_input = gr.Textbox(
-        label="C++ Code",
-        placeholder="e.g., 'int main() { int n; cin >> n; }'",
-        lines=5
-    )
-    generate_btn = gr.Button("Generate", variant="primary", elem_classes="btn-blue")
-    pseudocode_output = gr.Textbox(
-        label="Generated Pseudocode",
-        lines=5
-    )
-    
-    generate_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;
-}
-"""
+# 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()