Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,214 @@
+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 &nbsp;•&nbsp; Fine-tune: LoRA (r=8) &nbsp;•&nbsp; Trained on Kaggle P100
+        </div>
+        """
+    )
+
+if __name__ == "__main__":
+    demo.launch()