Update app.py

app.py

@@ -135,12 +135,12 @@ class Transformer(nn.Module):
 # Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load tokenizers (same files, but roles swapped)
+# 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
 
 # Load the model trained for C++ to pseudocode
-model_path = "c2p.pth"  # Assuming you retrained and saved as 'c2p.pth'
+model_path = "c2p.pth"  # Assuming retrained model for C++ to pseudocode
 model = torch.load(model_path, map_location=device, weights_only=False)
 model.eval()
 model = model.to(device)
@@ -149,7 +149,7 @@ model = model.to(device)
 def generate_pseudocode(cpp_code, max_len=500):
     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 (assuming ID=2)
+    tgt = torch.tensor([[2]], dtype=torch.long, device=device)  # <BOS> token (ID=2)
     
     generated_tokens = [2]  # Start with <BOS>
     with torch.no_grad():
@@ -158,7 +158,7 @@ def generate_pseudocode(cpp_code, max_len=500):
             next_token = output[:, -1, :].argmax(-1).item()
             generated_tokens.append(next_token)
             tgt = torch.cat([tgt, torch.tensor([[next_token]], device=device)], dim=1)
-            if next_token == 3:  # <EOS> token (assuming ID=3)
+            if next_token == 3:  # <EOS> token (ID=3)
                 break
     
     response = sp_pseudo.decode_ids(generated_tokens)  # Decode using pseudocode tokenizer
@@ -173,9 +173,16 @@ def generate_output(cpp_code):
 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.Code(label="C++ Code", language="cpp", placeholder="e.g., 'int main() { int n; cin >> n; }'")
+    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")
+    pseudocode_output = gr.Code(
+        label="Generated Pseudocode",
+        language="plaintext"  # Pseudocode isn’t a formal language, so use plaintext
+    )
     
     generate_btn.click(
         fn=generate_output,
@@ -185,7 +192,7 @@ with gr.Blocks(title="C++ to Pseudocode Transformer") as demo:
 
 demo.launch()
 
-# Custom CSS (unchanged)
+# Custom CSS
 demo.css = """
 .btn-blue {
     background-color: #007bff;