Update app.py

app.py

@@ -1,42 +1,27 @@
-import os
 import re
-import threading
 
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
-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"
+MODEL_ID = "shibsankardhara2/Qwen2.5-Coder-1.5B-Java-CSharp_V2"
 
-STAGE1_RESPONSE_MARKER = "### Response:\n\n"
-STAGE2_RESPONSE_MARKER = "### Response\n"
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# Prevent two users from running CPU inference simultaneously.
-generation_lock = threading.Lock()
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
 
-
-print("Downloading GGUF model...")
-
-model_path = hf_hub_download(
-    repo_id=MODEL_REPO,
-    filename=MODEL_FILE,
-)
-
-print("Loading GGUF model...")
-
-llm = Llama(
-    model_path=model_path,
-    n_ctx=2048,
-    n_threads=max(1, os.cpu_count() or 2),
-    n_threads_batch=max(1, os.cpu_count() or 2),
-    n_batch=64,
-    n_gpu_layers=0,
-    use_mmap=True,
-    verbose=False,
+print("Loading model on CPU...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
 )
 
+model.eval()
 print("Model loaded successfully.")
 
 
@@ -49,29 +34,29 @@ def add_java_hint(instruction: str) -> str:
     return f"{instruction} Write the solution in Java."
 
 
-def build_nl_to_java_prompt(instruction: str) -> str:
-    return (
-        "### Instruction:\n\n"
-        f"{add_java_hint(instruction)}\n\n"
-        "### Response:\n\n"
-    )
+def build_prompt(task: str, user_input: str) -> str:
+    user_input = user_input.strip()
 
+    if task == "Natural Language → Java":
+        return (
+            "### Instruction:\n\n"
+            f"{add_java_hint(user_input)}\n\n"
+            "### Response:\n\n"
+        )
 
-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"
+        f"{user_input}\n\n"
         "### Response\n"
     )
 
 
-def clean_generated_code(text: str) -> str:
+def clean_output(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,
@@ -99,54 +84,56 @@ def clean_generated_code(text: str) -> str:
     return text
 
 
-def run_generation(prompt: str, max_tokens: int) -> str:
-    with generation_lock:
-        response = llm(
-            prompt=prompt,
-            max_tokens=max_tokens,
-            temperature=0.0,
-            top_p=1.0,
-            repeat_penalty=1.0,
-            echo=False,
-            stop=[
-                "</s>",
-                "<|endoftext|>",
-                "<|im_end|>",
-                "### Instruction:",
-                "### Instruction\n",
-            ],
-        )
-
-    return response["choices"][0]["text"]
-
-
+@spaces.GPU(duration=120)
 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"
+    prompt = build_prompt(task, user_input)
 
-        else:
-            prompt = build_java_to_csharp_prompt(user_input)
-            generated = run_generation(prompt, max_tokens=400)
-            language = "csharp"
+    max_new_tokens = 300 if task == "Natural Language → Java" else 400
 
-        code = clean_generated_code(generated)
+    try:
+        model.to("cuda")
+
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=2048,
+        ).to("cuda")
+
+        with torch.inference_mode():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+
+        generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+
+        generated_text = tokenizer.decode(
+            generated_tokens,
+            skip_special_tokens=True,
+        )
+
+        code = clean_output(generated_text)
 
         if not code:
             return "The model returned an empty response. Please try again."
 
+        language = "java" if task == "Natural Language → Java" else "csharp"
+
         return f"```{language}\n{code}\n```"
 
     except Exception as error:
-        return (
-            "Generation failed.\n\n"
-            f"Error: {type(error).__name__}: {error}"
-        )
+        return f"Generation failed: {type(error).__name__}: {error}"
+
+    finally:
+        model.to("cpu")
+        torch.cuda.empty_cache()
 
 
 def update_input(task: str):
@@ -258,8 +245,4 @@ if __name__ == "__main__":
     demo.queue(
         default_concurrency_limit=1,
         max_size=10,
-    ).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-    )
+    ).launch()