Create app.py

app.py

@@ -1,4 +1,6 @@
 import os
+import re
+import threading
 
 import gradio as gr
 from huggingface_hub import hf_hub_download
@@ -8,6 +10,12 @@ from llama_cpp import Llama
 MODEL_REPO = "TaruniSwathi/Qwen2.5-Coder-1.5B-Java-CSharp-GGUF"
 MODEL_FILE = "Qwen2.5-Coder-1.5B-Java-CSharp_V2.Q4_K_M.gguf"
 
+STAGE1_RESPONSE_MARKER = "### Response:\n\n"
+STAGE2_RESPONSE_MARKER = "### Response\n"
+
+# Prevent two users from running CPU inference simultaneously.
+generation_lock = threading.Lock()
+
 
 print("Downloading GGUF model...")
 
@@ -16,14 +24,13 @@ model_path = hf_hub_download(
     filename=MODEL_FILE,
 )
 
-print("Model downloaded:", model_path)
-print("Loading model with llama.cpp...")
-
+print("Loading GGUF model...")
 
 llm = Llama(
     model_path=model_path,
-    n_ctx=2048,
+    n_ctx=4096,
     n_threads=max(1, os.cpu_count() or 2),
+    n_threads_batch=max(1, os.cpu_count() or 2),
     n_batch=128,
     n_gpu_layers=0,
     verbose=False,
@@ -32,84 +39,116 @@ llm = Llama(
 print("Model loaded successfully.")
 
 
-def build_prompt(task: str, user_input: str) -> str:
-    user_input = user_input.strip()
+def add_java_hint(instruction: str) -> str:
+    instruction = instruction.strip()
 
-    if task == "Natural Language → Java":
-        return f"""### Instruction:
-{user_input}
+    if "java" in instruction.lower():
+        return instruction
 
-### Java Code:
-"""
+    return f"{instruction} Write the solution in Java."
 
-    return f"""### Instruction:
-Translate the following Java code to C#.
 
-### Java Code:
-{user_input}
+def build_nl_to_java_prompt(instruction: str) -> str:
+    return (
+        "### Instruction:\n\n"
+        f"{add_java_hint(instruction)}\n\n"
+        "### Response:\n\n"
+    )
 
-### C# Code:
-"""
+
+def build_java_to_csharp_prompt(java_code: str) -> str:
+    return (
+        "### Instruction\n"
+        "Translate the following Java code into equivalent C#. "
+        "Write the solution in C#.\n\n"
+        "### Java\n"
+        f"{java_code.strip()}\n\n"
+        "### Response\n"
+    )
 
 
-def clean_output(task: str, generated_text: str) -> str:
-    generated_text = generated_text.strip()
+def clean_generated_code(text: str) -> str:
+    text = text.strip()
+
+    # Remove Markdown code fences if the model adds them.
+    fenced = re.search(
+        r"```(?:java|csharp|cs|c#)?\s*(.*?)```",
+        text,
+        flags=re.DOTALL | re.IGNORECASE,
+    )
+
+    if fenced:
+        text = fenced.group(1).strip()
 
     stop_markers = [
         "### Instruction:",
-        "### Java Code:",
-        "### C# Code:",
+        "### Instruction\n",
+        "### Java:",
+        "### Java\n",
+        "### Response:",
+        "### Response\n",
+        "<|im_start|>",
+        "<|im_end|>",
     ]
 
     for marker in stop_markers:
-        if marker in generated_text:
-            generated_text = generated_text.split(marker)[0].strip()
-
-    if task == "Natural Language → Java":
-        language = "java"
-    else:
-        language = "csharp"
-
-    return f"```{language}\n{generated_text}\n```"
+        if marker in text:
+            text = text.split(marker, 1)[0].strip()
 
+    return text
 
-def generate_code(
-    task: str,
-    user_input: str,
-    max_tokens: int,
-) -> str:
-    if not user_input or not user_input.strip():
-        return "Please enter a description or Java code."
-
-    prompt = build_prompt(task, user_input)
 
-    try:
+def run_generation(prompt: str, max_tokens: int) -> str:
+    with generation_lock:
         response = llm(
-            prompt,
-            max_tokens=int(max_tokens),
+            prompt=prompt,
+            max_tokens=max_tokens,
             temperature=0.0,
             top_p=1.0,
-            repeat_penalty=1.05,
+            repeat_penalty=1.0,
+            echo=False,
             stop=[
+                "</s>",
+                "<|endoftext|>",
+                "<|im_end|>",
                 "### Instruction:",
-                "### Java Code:",
-                "### C# Code:",
+                "### Instruction\n",
             ],
-            echo=False,
         )
 
-        generated_text = response["choices"][0]["text"]
+    return response["choices"][0]["text"]
+
+
+def generate_code(task: str, user_input: str) -> str:
+    if not user_input or not user_input.strip():
+        return "Please enter a requirement or Java code."
+
+    try:
+        if task == "Natural Language → Java":
+            prompt = build_nl_to_java_prompt(user_input)
+            generated = run_generation(prompt, max_tokens=300)
+            language = "java"
+
+        else:
+            prompt = build_java_to_csharp_prompt(user_input)
+            generated = run_generation(prompt, max_tokens=400)
+            language = "csharp"
+
+        code = clean_generated_code(generated)
 
-        if not generated_text.strip():
-            return "The model returned an empty response. Please try a more detailed input."
+        if not code:
+            return "The model returned an empty response. Please try again."
 
-        return clean_output(task, generated_text)
+        return f"```{language}\n{code}\n```"
 
     except Exception as error:
-        return f"Generation failed: {error}"
+        return (
+            "Generation failed.\n\n"
+            f"Error: {type(error).__name__}: {error}"
+        )
 
 
-def update_placeholder(task: str):
+def update_input(task: str):
     if task == "Natural Language → Java":
         return gr.update(
             label="Natural-language requirement",
@@ -126,7 +165,9 @@ def update_placeholder(task: str):
             "Example:\n"
             "public static int factorial(int n) {\n"
             "    int result = 1;\n"
-            "    for (int i = 2; i <= n; i++) result *= i;\n"
+            "    for (int i = 2; i <= n; i++) {\n"
+            "        result *= i;\n"
+            "    }\n"
             "    return result;\n"
             "}"
         ),
@@ -137,11 +178,11 @@ def update_placeholder(task: str):
 with gr.Blocks(title="Java and C# CodeGen") as demo:
     gr.Markdown(
         """
-        # Java and C# CodeGen
+# Java and C# CodeGen
 
-        Generate Java code from natural-language requirements or translate
-        Java code into C# using a fine-tuned Qwen2.5-Coder model.
-        """
+Generate Java code from natural-language requirements or translate Java code
+into equivalent C# using a fine-tuned Qwen2.5-Coder model.
+"""
     )
 
     task = gr.Dropdown(
@@ -159,15 +200,7 @@ with gr.Blocks(title="Java and C# CodeGen") as demo:
             "Example: Write a Java method to check whether "
             "a number is prime."
         ),
-        lines=12,
-    )
-
-    max_tokens = gr.Slider(
-        minimum=64,
-        maximum=512,
-        value=256,
-        step=32,
-        label="Maximum generated tokens",
+        lines=14,
     )
 
     generate_button = gr.Button(
@@ -175,12 +208,10 @@ with gr.Blocks(title="Java and C# CodeGen") as demo:
         variant="primary",
     )
 
-    output = gr.Markdown(
-        label="Generated code",
-    )
+    output = gr.Markdown()
 
     task.change(
-        fn=update_placeholder,
+        fn=update_input,
         inputs=task,
         outputs=user_input,
     )
@@ -190,7 +221,6 @@ with gr.Blocks(title="Java and C# CodeGen") as demo:
         inputs=[
             task,
             user_input,
-            max_tokens,
         ],
         outputs=output,
     )
@@ -200,12 +230,10 @@ with gr.Blocks(title="Java and C# CodeGen") as demo:
             [
                 "Natural Language → Java",
                 "Write a Java method to calculate factorial of a number using a loop.",
-                256,
             ],
             [
                 "Natural Language → Java",
                 "Write a Java method to reverse a string.",
-                256,
             ],
             [
                 "Java → C#",
@@ -216,16 +244,17 @@ with gr.Blocks(title="Java and C# CodeGen") as demo:
     }
     return result;
 }""",
-                256,
             ],
         ],
         inputs=[
             task,
             user_input,
-            max_tokens,
         ],
     )
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch()
+    demo.queue(
+        default_concurrency_limit=1,
+        max_size=10,
+    ).launch()