| import gradio as gr |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| import torch |
|
|
| MODEL_ID = "arinbalyan/code-translation-lora" |
|
|
| theme = ( |
| gr.themes.Soft(primary_hue="indigo", neutral_hue="slate") |
| .set(button_primary_background_fill_hover="#4f46e5") |
| ) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_ID, |
| torch_dtype=torch.float16, |
| device_map="auto", |
| ) |
|
|
| LANGUAGES = [ |
| "Python", "JavaScript", "TypeScript", "Java", "Go", |
| "Rust", "C++", "Ruby", "PHP", "C#", "Swift", "Kotlin", |
| ] |
|
|
|
|
| def generate( |
| task: str, |
| source_lang: str, |
| target_lang: str, |
| source_code: str, |
| description: str, |
| ) -> str: |
| if task == "Code Translation": |
| if not source_code.strip(): |
| return "# Enter source code to translate." |
| instruction = f"Translate the following {source_lang} code to {target_lang}:" |
| prompt = ( |
| f"{instruction}\n\n" |
| f"Source code ({source_lang}):\n" |
| f"```{source_lang.lower()}\n" |
| f"{source_code.strip()}\n" |
| f"```\n\n" |
| f"Translated code ({target_lang}):\n" |
| f"```{target_lang.lower()}\n" |
| ) |
| else: |
| if not description.strip(): |
| return "# Describe what you want to code." |
| prompt = ( |
| f"### Instruction:\n{description.strip()}\n\n" |
| f"### Code:\n" |
| ) |
|
|
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) |
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=512, |
| temperature=0.3, |
| do_sample=True, |
| pad_token_id=tokenizer.eos_token_id, |
| ) |
| result = tokenizer.decode(outputs[0], skip_special_tokens=True) |
| code = result[len(prompt):] |
| if code.endswith("```"): |
| code = code[:-3].strip() |
| return code.strip() |
|
|
|
|
| def update_visibility(task: str): |
| return { |
| translation_box: gr.update(visible=task == "Code Translation"), |
| translation_source_lang: gr.update(visible=task == "Code Translation"), |
| translation_target_lang: gr.update(visible=task == "Code Translation"), |
| translation_target_label: gr.update(visible=task == "Code Translation"), |
| generation_box: gr.update(visible=task == "Code Generation"), |
| generation_label: gr.update(visible=task == "Code Generation"), |
| } |
|
|
|
|
| with gr.Blocks(theme=theme, title="Code Translation") as demo: |
| gr.Markdown( |
| """ |
| # 🔄 Code Translation Studio |
| Fine-tuned Qwen2.5-Coder-1.5B on code translation and generation tasks. |
| Translate code between languages or generate code from descriptions. |
| """ |
| ) |
|
|
| task = gr.Radio( |
| choices=["Code Translation", "Code Generation"], |
| value="Code Translation", |
| label="Task", |
| ) |
|
|
| with gr.Row(): |
| with gr.Column(scale=1): |
| translation_box = gr.Textbox( |
| label="Source Code", |
| placeholder="Paste your source code here...", |
| lines=10, |
| visible=True, |
| ) |
| translation_source_lang = gr.Dropdown( |
| choices=LANGUAGES, |
| value="Python", |
| label="Source Language", |
| visible=True, |
| ) |
| translation_target_lang = gr.Dropdown( |
| choices=LANGUAGES, |
| value="JavaScript", |
| label="Target Language", |
| visible=True, |
| ) |
|
|
| with gr.Column(scale=1): |
| translation_target_label = gr.Code( |
| label="Translated Code", |
| language="javascript", |
| lines=10, |
| visible=True, |
| ) |
|
|
| generation_box = gr.Textbox( |
| label="Description", |
| placeholder="e.g., Write a function to merge two sorted lists in Python...", |
| lines=4, |
| visible=False, |
| ) |
| generation_label = gr.Code( |
| label="Generated Code", |
| language="python", |
| lines=10, |
| visible=False, |
| ) |
|
|
| run_btn = gr.Button("Generate", variant="primary", size="lg") |
|
|
| gr.Examples( |
| examples=[ |
| [ |
| "Code Translation", |
| "def add(a, b):\n return a + b", |
| "Python", |
| "JavaScript", |
| "", |
| ], |
| [ |
| "Code Translation", |
| "function isEven(n) { return n % 2 === 0; }", |
| "JavaScript", |
| "Python", |
| "", |
| ], |
| [ |
| "Code Generation", |
| "", |
| "Python", |
| "JavaScript", |
| "Write a Python function to reverse a string.", |
| ], |
| [ |
| "Code Generation", |
| "", |
| "Python", |
| "JavaScript", |
| "Write a function to check if a number is prime in Python.", |
| ], |
| ], |
| inputs=[ |
| task, |
| translation_box, |
| translation_source_lang, |
| translation_target_lang, |
| generation_box, |
| ], |
| label="Try one of these", |
| ) |
|
|
| task.change( |
| fn=update_visibility, |
| inputs=task, |
| outputs=[ |
| translation_box, |
| translation_source_lang, |
| translation_target_lang, |
| translation_target_label, |
| generation_box, |
| generation_label, |
| ], |
| ) |
|
|
| run_btn.click( |
| fn=generate, |
| inputs=[ |
| task, |
| translation_source_lang, |
| translation_target_lang, |
| translation_box, |
| generation_box, |
| ], |
| outputs=[translation_target_label, generation_label], |
| ) |
|
|
| gr.Markdown( |
| """ |
| <div style="text-align:center; margin-top:1rem; color:gray; font-size:0.9rem;"> |
| Base: Qwen2.5-Coder-1.5B • Fine-tune: LoRA (r=8) • Trained on Kaggle P100 |
| </div> |
| """ |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
