Update app.py

app.py

@@ -19,14 +19,17 @@ MODEL_ID = "shibsankardhara2/Qwen2.5-Coder-1.5B-Java-CSharp_V5"
 
 print("Loading tokenizer...")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
 
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
 
 # ============================================================
-# LOAD MODEL
+# LOAD MODEL ON CPU
 # ============================================================
 
 print("Loading model on CPU...")
@@ -35,6 +38,7 @@ model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
     low_cpu_mem_usage=True,
+    trust_remote_code=True,
 )
 
 model.eval()
@@ -48,8 +52,8 @@ print("Model loaded successfully.")
 
 def add_java_hint(instruction: str) -> str:
     """
-    Add an explicit Java instruction when the user has not
-    mentioned Java in the natural-language requirement.
+    Explicitly mention Java when it is missing from the
+    natural-language instruction.
     """
     instruction = instruction.strip()
 
@@ -61,7 +65,7 @@ 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.
+    Construct the prompt expected by the fine-tuned model.
     """
     user_input = user_input.strip()
 
@@ -83,25 +87,30 @@ def build_prompt(task: str, user_input: str) -> str:
 
 
 # ============================================================
-# OUTPUT CLEANING
+# GENERAL OUTPUT CLEANING
 # ============================================================
 
-def clean_output(text: str) -> str:
+def clean_output(generated_text: str) -> str:
     """
-    Remove Markdown code fences, repeated prompt sections,
-    and model-specific special tokens from the generated output.
+    Remove Markdown fences, repeated prompt sections and
+    model-specific special tokens.
     """
-    text = text.strip()
+    if not generated_text:
+        return ""
+
+    cleaned_text = generated_text.strip()
 
-    fenced = re.search(
+    # Extract the contents of a Markdown code block, when present.
+    fenced_match = re.search(
         r"```(?:java|csharp|cs|c#)?\s*(.*?)```",
-        text,
+        cleaned_text,
         flags=re.DOTALL | re.IGNORECASE,
     )
 
-    if fenced:
-        text = fenced.group(1).strip()
+    if fenced_match:
+        cleaned_text = fenced_match.group(1).strip()
 
+    # Remove prompt sections that the model may generate again.
     stop_markers = [
         "### Instruction:",
         "### Instruction\n",
@@ -111,24 +120,49 @@ def clean_output(text: str) -> str:
         "### Response\n",
         "<|im_start|>",
         "<|im_end|>",
+        "<|endoftext|>",
     ]
 
     for marker in stop_markers:
-        if marker in text:
-            text = text.split(marker, 1)[0].strip()
+        marker_position = cleaned_text.find(marker)
+
+        if marker_position != -1:
+            cleaned_text = cleaned_text[:marker_position].strip()
+
+    # Remove remaining opening or closing fences.
+    cleaned_text = re.sub(
+        r"^```(?:java|csharp|cs|c#)?\s*",
+        "",
+        cleaned_text,
+        flags=re.IGNORECASE,
+    )
+
+    cleaned_text = re.sub(
+        r"\s*```$",
+        "",
+        cleaned_text,
+    )
+
+    return cleaned_text.strip()
 
-    return text
 
+# ============================================================
+# C# OUTPUT CLEANING
+# ============================================================
 
-def clean_csharp_output(code: str) -> str:
+def clean_csharp_output(text: str) -> str:
     """
-    Remove unnecessary `virtual` and `override` modifiers from
-    standalone C# method snippets.
+    Remove invalid virtual or 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 the output contains a class, struct, interface, record or
+    enum declaration, modifiers are preserved.
     """
-    code = code.strip()
+    if not text:
+        return ""
+
+    # Explicitly initialise code before it is accessed.
+    code = text.strip()
 
     has_type_declaration = re.search(
         r"\b(class|struct|interface|record|enum)\b",
@@ -139,8 +173,8 @@ def clean_csharp_output(code: str) -> str:
     if has_type_declaration:
         return code
 
-    # Example:
-    # public virtual int Add(...) -> public int Add(...)
+    # public virtual int Method() -> public int Method()
+    # protected override void Method() -> protected void Method()
     code = re.sub(
         r"\b(public|private|protected|internal)\s+"
         r"(?:virtual|override)\s+",
@@ -149,8 +183,8 @@ def clean_csharp_output(code: str) -> str:
         flags=re.IGNORECASE,
     )
 
-    # Example:
-    # virtual int Add(...) -> int Add(...)
+    # virtual int Method() -> int Method()
+    # override void Method() -> void Method()
     code = re.sub(
         r"(^|\n)(\s*)(?:virtual|override)\s+",
         r"\1\2",
@@ -168,66 +202,80 @@ def clean_csharp_output(code: str) -> str:
 @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#.
+    Generate Java from natural language or translate Java to C#.
     """
     if not user_input or not user_input.strip():
         return "Please enter a requirement or Java code."
 
     prompt = build_prompt(task, user_input)
 
-    if task == "Natural Language → Java":
-        max_new_tokens = 300
-    else:
-        max_new_tokens = 400
+    max_new_tokens = (
+        300
+        if task == "Natural Language → Java"
+        else 400
+    )
 
     try:
-        # Move the model to the GPU only during generation.
+        if not torch.cuda.is_available():
+            return "Generation failed: GPU is not available."
+
+        # Move the model to the allocated ZeroGPU device.
         model.to("cuda")
+        model.eval()
 
-        inputs = tokenizer(
+        tokenized_inputs = tokenizer(
             prompt,
             return_tensors="pt",
             truncation=True,
             max_length=2048,
-        ).to("cuda")
+        )
+
+        tokenized_inputs = {
+            key: value.to("cuda")
+            for key, value in tokenized_inputs.items()
+        }
 
         with torch.inference_mode():
-            outputs = model.generate(
-                **inputs,
+            generated_ids = model.generate(
+                **tokenized_inputs,
                 max_new_tokens=max_new_tokens,
                 do_sample=False,
+                use_cache=True,
                 pad_token_id=tokenizer.pad_token_id,
                 eos_token_id=tokenizer.eos_token_id,
             )
 
-        # Remove the original input prompt from the generated sequence.
-        prompt_length = inputs["input_ids"].shape[1]
+        prompt_length = tokenized_inputs["input_ids"].shape[1]
 
-        generated_tokens = outputs[0][prompt_length:]
+        new_token_ids = generated_ids[0][prompt_length:]
 
         generated_text = tokenizer.decode(
-            generated_tokens,
+            new_token_ids,
             skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
         )
 
-        # Clean repeated prompts, code fences, and special tokens.
-        code = clean_output(generated_text)
+        # Always run the general output cleaner.
+        output_code = clean_output(generated_text)
 
-        # Apply the C# modifier cleanup only for Java → C#.
+        # Run the C#-specific cleaner only for Java → C#.
         if task == "Java → C#":
-            code = clean_csharp_output(code)
-            language = "csharp"
+            output_code = clean_csharp_output(output_code)
+            markdown_language = "csharp"
         else:
-            language = "java"
+            markdown_language = "java"
 
-        if not code:
+        if not output_code:
             return (
                 "The model returned an empty response. "
-                "Please try again."
+                "Please try a more specific input."
             )
 
-        return f"```{language}\n{code}\n```"
+        return (
+            f"```{markdown_language}\n"
+            f"{output_code}\n"
+            "```"
+        )
 
     except Exception as error:
         return (
@@ -236,21 +284,23 @@ def generate_code(task: str, user_input: str) -> str:
         )
 
     finally:
-        # Move the model back to CPU to release the ZeroGPU allocation.
-        model.to("cpu")
+        # Return the model to CPU after the ZeroGPU request.
+        try:
+            model.to("cpu")
+        except Exception as move_error:
+            print(f"Could not move model to CPU: {move_error}")
 
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
 
 # ============================================================
-# GRADIO INPUT UPDATE
+# UPDATE TEXTBOX
 # ============================================================
 
 def update_input(task: str):
     """
-    Change the textbox label and placeholder based on the
-    selected task.
+    Update the input textbox for the selected task.
     """
     if task == "Natural Language → Java":
         return gr.update(