Update app.py

app.py

@@ -1,26 +1,63 @@
 """
-Gradio demo for Gemma Code Generator using HuggingFace Inference API.
-This runs serverless on HF infrastructure - no GPU costs!
+Gradio demo for Gemma Code Generator.
+Loads the fine-tuned model directly using PEFT.
 """
 
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 
 # Model configuration
-MODEL_NAME = "nvhuynh16/gemma-2b-code-alpaca-best"  # Best checkpoint (step 2000)
+BASE_MODEL = "google/gemma-2-2b-it"
+ADAPTER_MODEL = "nvhuynh16/gemma-2b-code-alpaca-best"
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+# Global variables for lazy loading
+tokenizer = None
+model = None
+
+def load_model():
+    """Lazy load model on first request"""
+    global tokenizer, model
+
+    if model is None:
+        print("Loading model for the first time...")
+
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)
+
+        # Load base model with 4-bit quantization
+        base_model = AutoModelForCausalLM.from_pretrained(
+            BASE_MODEL,
+            device_map="auto",
+            torch_dtype=torch.float16,
+            load_in_4bit=True,
+            token=HF_TOKEN
+        )
+
+        # Load LoRA adapter
+        model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL, token=HF_TOKEN)
+        model.eval()
+
+        print("Model loaded successfully!")
 
-# Initialize Inference client (no model specified here - we'll pass it per request)
-client = InferenceClient()
+    return tokenizer, model
 
 
 def generate_code(instruction: str, max_tokens: int = 256, temperature: float = 0.7):
-    """Generate code from instruction using HF Inference API"""
+    """Generate code from instruction"""
 
     if not instruction.strip():
         return "Please enter an instruction."
 
-    # Format prompt in Alpaca style
-    prompt = f"""### Instruction:
+    try:
+        # Load model (cached after first call)
+        tok, mdl = load_model()
+
+        # Format prompt in Alpaca style
+        prompt = f"""### Instruction:
 {instruction}
 
 ### Input:
@@ -29,32 +66,33 @@ def generate_code(instruction: str, max_tokens: int = 256, temperature: float =
 ### Response:
 """
 
-    try:
-        # Generate using HF Inference API
-        response = client.text_generation(
-            prompt,
-            model=MODEL_NAME,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=0.9,
-            do_sample=True,
-            return_full_text=False,
-        )
+        # Tokenize
+        inputs = tok(prompt, return_tensors="pt").to(mdl.device)
+
+        # Generate
+        with torch.no_grad():
+            outputs = mdl.generate(
+                **inputs,
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tok.eos_token_id,
+            )
 
-        return response.strip()
+        # Decode
+        generated = tok.decode(outputs[0], skip_special_tokens=True)
 
-    except Exception as e:
-        error_msg = str(e)
-        if "410" in error_msg or "Gone" in error_msg:
-            return "⚠️ API endpoint error. This usually means the Inference API is updating. Please try again in a moment."
-        elif "Model too large" in error_msg or "not currently loaded" in error_msg or "loading" in error_msg.lower():
-            return "⏳ Model is loading (first request takes 1-2 minutes). Please try again in a moment."
-        elif "rate limit" in error_msg.lower():
-            return "⚠️ Rate limit reached. Please wait a few minutes and try again."
-        elif "404" in error_msg or "not found" in error_msg.lower():
-            return "⚠️ Model not found or not enabled for Inference API. Please check the model settings on HuggingFace."
+        # Extract code after "### Response:"
+        if "### Response:" in generated:
+            code = generated.split("### Response:")[-1].strip()
         else:
-            return f"Error: {error_msg}\n\nPlease try again. If the issue persists, the model may be loading for the first time."
+            code = generated.strip()
+
+        return code
+
+    except Exception as e:
+        return f"Error: {str(e)}\n\nPlease try again."
 
 
 # Custom CSS for better appearance