Update app.py

app.py

@@ -6,15 +6,31 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
+# ============================================================
+# MODEL CONFIGURATION
+# ============================================================
+
 MODEL_ID = "shibsankardhara2/Qwen2.5-Coder-1.5B-Java-CSharp_V5"
 
+
+# ============================================================
+# LOAD TOKENIZER
+# ============================================================
+
 print("Loading tokenizer...")
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
+
+# ============================================================
+# LOAD MODEL
+# ============================================================
+
 print("Loading model on CPU...")
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
@@ -22,10 +38,19 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 model.eval()
+
 print("Model loaded successfully.")
 
 
+# ============================================================
+# PROMPT FUNCTIONS
+# ============================================================
+
 def add_java_hint(instruction: str) -> str:
+    """
+    Add an explicit Java instruction when the user has not
+    mentioned Java in the natural-language requirement.
+    """
     instruction = instruction.strip()
 
     if "java" in instruction.lower():
@@ -35,6 +60,9 @@ def add_java_hint(instruction: str) -> str:
 
 
 def build_prompt(task: str, user_input: str) -> str:
+    """
+    Build the appropriate prompt based on the selected task.
+    """
     user_input = user_input.strip()
 
     if task == "Natural Language → Java":
@@ -45,16 +73,24 @@ def build_prompt(task: str, user_input: str) -> str:
         )
 
     return (
-        "### Instruction\n"
+        "### Instruction:\n\n"
         "Translate the following Java code into equivalent C#. "
-        "Write the solution in C#.\n\n"
-        "### Java\n"
+        "Write only the C# solution.\n\n"
+        "### Java:\n\n"
         f"{user_input}\n\n"
-        "### Response\n"
+        "### Response:\n\n"
     )
 
 
+# ============================================================
+# OUTPUT CLEANING
+# ============================================================
+
 def clean_output(text: str) -> str:
+    """
+    Remove Markdown code fences, repeated prompt sections,
+    and model-specific special tokens from the generated output.
+    """
     text = text.strip()
 
     fenced = re.search(
@@ -83,39 +119,70 @@ def clean_output(text: str) -> str:
 
     return text
 
-def _clean_csharp_output(text: str) -> str:
+
+def clean_csharp_output(code: str) -> str:
     """
-    Remove spurious `virtual` / `override` modifiers the model adds to bare
-    (class-less) method snippets.
-
-    The model was trained on C# methods that usually live inside a class, where
-    `public virtual ...` is common. When it translates a *standalone* Java method
-    it carries the `virtual` keyword over — but `virtual`/`override` are only
-    valid on members of a class, so on a bare snippet they are invalid C#.
-    We therefore strip them ONLY when the snippet has no enclosing type
-    declaration (class / struct / interface / record / enum).
+    Remove unnecessary `virtual` and `override` modifiers from
+    standalone C# method snippets.
+
+    These modifiers are preserved when the generated code contains
+    a class, struct, interface, record, or enum declaration.
     """
-    if re.search(r"\b(class|struct|interface|record|enum)\b", code):
-        return code  # real type present — leave modifiers intact
-    # Drop 'virtual'/'override' after an access modifier: 'public virtual int' -> 'public int'.
-    code = re.sub(r"\b(public|private|protected|internal)\s+(?:virtual|override)\s+",
-                  r"\1 ", code)
-    # Drop a leading 'virtual'/'override' with no access modifier.
-    code = re.sub(r"(^|\n)(\s*)(?:virtual|override)\s+", r"\1\2", code)
-    return code
+    code = code.strip()
+
+    has_type_declaration = re.search(
+        r"\b(class|struct|interface|record|enum)\b",
+        code,
+        flags=re.IGNORECASE,
+    )
+
+    if has_type_declaration:
+        return code
+
+    # Example:
+    # public virtual int Add(...) -> public int Add(...)
+    code = re.sub(
+        r"\b(public|private|protected|internal)\s+"
+        r"(?:virtual|override)\s+",
+        r"\1 ",
+        code,
+        flags=re.IGNORECASE,
+    )
+
+    # Example:
+    # virtual int Add(...) -> int Add(...)
+    code = re.sub(
+        r"(^|\n)(\s*)(?:virtual|override)\s+",
+        r"\1\2",
+        code,
+        flags=re.IGNORECASE,
+    )
 
+    return code.strip()
+
+
+# ============================================================
+# CODE GENERATION
+# ============================================================
 
-    
 @spaces.GPU(duration=120)
 def generate_code(task: str, user_input: str) -> str:
+    """
+    Generate Java code from natural language or translate
+    Java code into C#.
+    """
     if not user_input or not user_input.strip():
         return "Please enter a requirement or Java code."
 
     prompt = build_prompt(task, user_input)
 
-    max_new_tokens = 300 if task == "Natural Language → Java" else 400
+    if task == "Natural Language → Java":
+        max_new_tokens = 300
+    else:
+        max_new_tokens = 400
 
     try:
+        # Move the model to the GPU only during generation.
         model.to("cuda")
 
         inputs = tokenizer(
@@ -130,35 +197,61 @@ def generate_code(task: str, user_input: str) -> str:
                 **inputs,
                 max_new_tokens=max_new_tokens,
                 do_sample=False,
-                pad_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.pad_token_id,
                 eos_token_id=tokenizer.eos_token_id,
             )
 
-        generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+        # Remove the original input prompt from the generated sequence.
+        prompt_length = inputs["input_ids"].shape[1]
+
+        generated_tokens = outputs[0][prompt_length:]
 
         generated_text = tokenizer.decode(
             generated_tokens,
             skip_special_tokens=True,
         )
 
-        code_ = _clean_csharp_output(generated_text)
+        # Clean repeated prompts, code fences, and special tokens.
+        code = clean_output(generated_text)
 
-        if not code_:
-            return "The model returned an empty response. Please try again."
+        # Apply the C# modifier cleanup only for Java → C#.
+        if task == "Java → C#":
+            code = clean_csharp_output(code)
+            language = "csharp"
+        else:
+            language = "java"
 
-        language = "java" if task == "Natural Language → Java" else "csharp"
+        if not code:
+            return (
+                "The model returned an empty response. "
+                "Please try again."
+            )
 
-        return f"```{language}\n{code_}\n```"
+        return f"```{language}\n{code}\n```"
 
     except Exception as error:
-        return f"Generation failed: {type(error).__name__}: {error}"
+        return (
+            "Generation failed: "
+            f"{type(error).__name__}: {error}"
+        )
 
     finally:
+        # Move the model back to CPU to release the ZeroGPU allocation.
         model.to("cpu")
-        torch.cuda.empty_cache()
 
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+
+# ============================================================
+# GRADIO INPUT UPDATE
+# ============================================================
 
 def update_input(task: str):
+    """
+    Change the textbox label and placeholder based on the
+    selected task.
+    """
     if task == "Natural Language → Java":
         return gr.update(
             label="Natural-language requirement",
@@ -185,6 +278,10 @@ def update_input(task: str):
     )
 
 
+# ============================================================
+# GRADIO INTERFACE
+# ============================================================
+
 with gr.Blocks(title="Java and C# CodeGen") as demo:
     gr.Markdown(
         """
@@ -218,7 +315,9 @@ into equivalent C# using a fine-tuned Qwen2.5-Coder model.
         variant="primary",
     )
 
-    output = gr.Markdown()
+    output = gr.Markdown(
+        value="Generated code will appear here."
+    )
 
     task.change(
         fn=update_input,
@@ -239,20 +338,38 @@ into equivalent C# using a fine-tuned Qwen2.5-Coder model.
         examples=[
             [
                 "Natural Language → Java",
-                "Write a Java method to calculate factorial of a number using a loop.",
+                (
+                    "Write a Java method to calculate factorial "
+                    "of a number using a loop."
+                ),
             ],
             [
                 "Natural Language → Java",
                 "Write a Java method to reverse a string.",
             ],
+            [
+                "Natural Language → Java",
+                (
+                    "Write a Java method to check whether "
+                    "a number is prime."
+                ),
+            ],
             [
                 "Java → C#",
                 """public static int factorial(int n) {
     int result = 1;
+
     for (int i = 2; i <= n; i++) {
         result *= i;
     }
+
     return result;
+}""",
+            ],
+            [
+                "Java → C#",
+                """public boolean isEven(int n) {
+    return n % 2 == 0;
 }""",
             ],
         ],
@@ -263,6 +380,9 @@ into equivalent C# using a fine-tuned Qwen2.5-Coder model.
     )
 
 
+# ============================================================
+# START APPLICATION
+# ============================================================
 
 if __name__ == "__main__":
     demo.queue(