Update app.py

app.py

@@ -1,67 +1,42 @@
 import gradio as gr
-from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
 
-# --- CONFIGURATION ---
-# The base model 
-BASE_MODEL = "unsloth/llama-3-8b" 
-# Fine-tuned adapter
-ADAPTER_MODEL = "nihardon/fine-tuned-unit-test-generator" 
-
-print(f"Loading {ADAPTER_MODEL} on CPU... (This might take a minute)")
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-
-# Load model (CPU optimized loading)
-model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    device_map="cpu",
-    torch_dtype=torch.float32, 
-    low_cpu_mem_usage=True
+# 1. Download your specific GGUF file
+model_path = hf_hub_download(
+    repo_id="nihardon/fine-tuned-unit-test-generator",
+    filename="llama-3-8b.Q4_K_M.gguf",  # <--- DOUBLE CHECK THIS NAME!
 )
 
-# Apply fine-tuned adapters
-model = PeftModel.from_pretrained(model, ADAPTER_MODEL)
+# 2. Load the Model (CPU Optimized)
+llm = Llama(
+    model_path=model_path,
+    n_ctx=2048,
+    n_threads=2, 
+)
 
 def generate_test(user_code):
-    alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+    prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
 
 ### Instruction:
 You are an expert Python QA engineer. Write a pytest unit test for the following function.
 
 ### Input:
-{}
+{user_code}
 
 ### Response:
 """
-    # Format and tokenize
-    prompt = alpaca_prompt.format(user_code)
-    inputs = tokenizer(prompt, return_tensors="pt")
-    
-    # Generate
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs, 
-            max_new_tokens=256, 
-            use_cache=True,
-            temperature=0.1 
-        )
-    
-    # Decode
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response.split("### Response:")[-1].strip()
+    output = llm(prompt, max_tokens=256, stop=["### Instruction:"], echo=False)
+    return output['choices'][0]['text'].strip()
 
-# UI
+# 3. The UI
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🧪 AI Unit Test Generator")
-    gr.Markdown(f"**Model:** {ADAPTER_MODEL} (Llama-3 Fine-Tune) | **Status:** Running on CPU")
-    gr.Markdown("Paste your Python function below, and the AI will write a Pytest case for it.")
+    gr.Markdown("**Model:** Custom Fine-Tuned Llama-3 (GGUF) | **Status:** Running Locally")
     
     with gr.Row():
         with gr.Column():
-            input_box = gr.Code(language="python", label="Paste Python Function Here", value="def add(a, b):\n    return a + b")
+            input_box = gr.Code(language="python", value="def add(a, b):\n    return a + b", label="Function")
             btn = gr.Button("Generate Pytest", variant="primary")
         with gr.Column():
             output_box = gr.Code(language="python", label="Generated Test Case")