Update app.py

app.py

@@ -140,34 +140,42 @@ sp_code = spm.SentencePieceProcessor(model_file="code_tokenizer.model")      # C
 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 retrained model for C++ to pseudocode
+model_path = "c2p.pth"  # Ensure this is the correct model for C++ to pseudocode
 model = torch.load(model_path, map_location=device, weights_only=False)
 model.eval()
 model = model.to(device)
 
-# Function to generate pseudocode from C++ code
+# Function to generate pseudocode from C++ code with streaming
 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 (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
+    
     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)
-            if next_token == 3:  # <EOS> token (ID=3)
+            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.")
                 break
+        print("Generation complete or max length reached.")
     
-    response = sp_pseudo.decode_ids(generated_tokens)  # Decode using pseudocode tokenizer
-    return response
+    yield response  # Final output
 
-# Gradio interface function
+# Gradio interface function with streaming
 def generate_output(cpp_code):
-    pseudocode = generate_pseudocode(cpp_code)
-    return pseudocode
+    for response in generate_pseudocode(cpp_code, max_len=500):
+        yield response
 
 # Gradio UI setup
 with gr.Blocks(title="C++ to Pseudocode Transformer") as demo:
@@ -187,7 +195,8 @@ with gr.Blocks(title="C++ to Pseudocode Transformer") as demo:
     generate_btn.click(
         fn=generate_output,
         inputs=[cpp_input],
-        outputs=pseudocode_output
+        outputs=pseudocode_output,
+        _js="() => [document.querySelector('#cpp_input textarea').value]"  # Ensure input is passed correctly
     )
 
 demo.launch()