iffsuleman/pseudocode-to-code

1import torch2import gradio as gr3from transformers import GPT2Tokenizer, GPT2LMHeadModel4from peft import PeftModel, PeftConfig5 6class PseudocodeToPython:7    def __init__(self, model_path="./pseudo-to-python-final"):8        try:9            config = PeftConfig.from_pretrained(model_path)10            self.tokenizer = GPT2Tokenizer.from_pretrained(config.base_model_name_or_path)11            self.tokenizer.pad_token = self.tokenizer.eos_token12            self.tokenizer.add_special_tokens({'additional_special_tokens': ['<|pseudo|>', '<|python|>', '<|end|>']})13 14            base_model = GPT2LMHeadModel.from_pretrained(config.base_model_name_or_path)15            base_model.resize_token_embeddings(len(self.tokenizer))16 17            self.model = PeftModel.from_pretrained(base_model, model_path)18            self.model.eval()19 20            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")21            self.model = self.model.to(self.device)22 23            print("Model loaded successfully")24 25        except Exception as e:26            print(f"Error: {e}")27            self.setup_fallback()28 29    def setup_fallback(self):30        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')31        self.tokenizer.pad_token = self.tokenizer.eos_token32        self.model = GPT2LMHeadModel.from_pretrained('gpt2')33        self.model.eval()34        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")35        self.model = self.model.to(self.device)36 37    def generate(self, pseudocode):38        input_text = f"<|pseudo|>{pseudocode}<|python|>"39        inputs = self.tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512).to(self.device)40 41        with torch.no_grad():42            outputs = self.model.generate(43                **inputs,44                max_new_tokens=256,45                num_beams=3,46                temperature=0.7,47                do_sample=True,48                pad_token_id=self.tokenizer.eos_token_id,49                early_stopping=True50            )51 52        full_output = self.tokenizer.decode(outputs[0], skip_special_tokens=False)53 54        if "<|python|>" in full_output:55            python_code = full_output.split("<|python|>")[1].strip()56            if "<|end|>" in python_code:57                python_code = python_code.split("<|end|>")[0].strip()58            return python_code59 60        return full_output.replace(input_text, "").strip()61 62def create_interface():63    generator = PseudocodeToPython()64 65    def generate_code(pseudocode):66        if not pseudocode.strip():67            return "Please enter pseudocode"68        try:69            return generator.generate(pseudocode)70        except Exception as e:71            return f"Error: {str(e)}"72 73    with gr.Blocks(title="Pseudocode to Python") as interface:74        gr.Markdown("# Pseudocode to Python Generator")75 76        with gr.Row():77            with gr.Column():78                input_text = gr.Textbox(79                    label="Pseudocode Input",80                    placeholder="Enter your pseudocode here...",81                    lines=482                )83                generate_btn = gr.Button("Generate Python Code", variant="primary")84 85            with gr.Column():86                output_text = gr.Textbox(87                    label="Generated Python Code",88                    placeholder="Python code will appear here...",89                    lines=490                )91 92        examples = [93            "print numbers from 1 to 10",94            "calculate sum of two numbers",95            "find maximum number in list"96        ]97 98        gr.Examples(examples=examples, inputs=input_text)99 100        generate_btn.click(fn=generate_code, inputs=input_text, outputs=output_text)101 102    return interface103 104interface = create_interface()105interface.launch(share=True)