Spaces:
Sleeping
Sleeping
File size: 1,705 Bytes
9b2f41b fd564a5 9b3eb0b 9b2f41b 3a24c67 4d5ef83 3c2266c 9b2f41b 3c2266c 0a17c6b 9b2f41b 3c2266c 3a24c67 3c2266c 9b2f41b fc2bdb8 3c2266c fc2bdb8 3c2266c fc2bdb8 12724c3 0a17c6b 3c2266c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import gradio as gr
import torch
import torch.nn as nn
import sentencepiece as spm
import math
# 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)
# Load the full saved model (architecture + weights)
model_path = "code2pseudo.pth"
model = torch.load(model_path, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"), weights_only=False)
model.eval()
# Function to generate pseudocode
def generate_pseudocode(cpp_code, max_len):
model.eval()
src = torch.tensor([sp_code.encode_as_ids(cpp_code)], dtype=torch.long) # Tokenize C++ code
tgt = torch.tensor([[2]], dtype=torch.long) # <bos_id>=2
generated_tokens = [2] # Start with <START>
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]])], dim=1)
if next_token == 3: # <END>=3
break
return sp_pseudo.decode_ids(generated_tokens) # Final decoded output
# Gradio interface
demo = gr.Interface(
fn=generate_pseudocode,
inputs=[
gr.Textbox(placeholder="Enter C++ code here", label="C++ Code"),
gr.Slider(minimum=10, maximum=1000, value=50, step=1, label="Max Tokens")
],
outputs=gr.Textbox(label="Generated Pseudocode"),
title="C++ to Pseudocode Converter",
description="Enter C++ code and get its pseudocode equivalent using a transformer model."
)
if __name__ == "__main__":
demo.launch()
|