Update app.py

app.py

@@ -18,23 +18,36 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # ---------------------------------------------------------------------------
 # MODULE 4 — Generation Module (LLM setup)
 # ---------------------------------------------------------------------------
-MODEL_ID = "google/gemma-2b-it-tflite"   # swap for "google/gemma-4-9b-it" if VRAM allows
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,          # efficient on modern GPUs/CPUs
-    device_map="auto",                    # auto-detect GPU/CPU
-)
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=512,
-    do_sample=False,                      # deterministic output
-)
-
-
+MODEL_ID = "google/gemma-4-it"   # swap for "google/gemma-4-9b-it" if VRAM allows
+
+_tokenizer=None
+_pipe=None
+
+def get_tokenizer():
+    """Returns the tokenizer, loading it on first call."""
+    global _tokenizer
+    if _tokenizer is None:
+        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    return _tokenizer
+ 
+def get_pipeline():
+    """Returns the generation pipeline, loading model on first call (lazy load)."""
+    global _pipe
+    if _pipe is None:
+        tokenizer = get_tokenizer()
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.bfloat16,   # efficient on modern GPUs/CPUs
+            device_map="auto",             # auto-detect GPU/CPU
+        )
+        _pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=512,
+            do_sample=False,               # deterministic output
+        )
+    return _pipe
 
 # ---------------------------------------------------------------------------
 # MODULE 3 — Pre-processing Module (6-Element Framework)
@@ -182,6 +195,13 @@ def validate_and_evaluate(description: str, code: str):
     # --- Pre-processing ---
     prompt = build_prompt(description, code)
 
+    token_count = len(get_tokenizer().encode(prompt))
+    if token_count > 2048:
+        return "", "", "", (
+            f"⚠️ Input too long ({token_count} tokens). "
+            "Please shorten your description or code and try again."
+        )
+    
     # --- Generation ---
     try:
         raw_output = generate(prompt)