Update app.py

app.py

@@ -2,14 +2,13 @@
 # Imports
 # ===============================
 import torch
-import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 # ===============================
-# Load Model & Tokenizer (ONCE)
+# Load Model & Tokenizer (CPU)
 # ===============================
-print("Loading Aksara v1 model...")
+print("Loading Aksara v1 model (CPU)...")
 
 tokenizer = AutoTokenizer.from_pretrained(
     "cropinailab/aksara_v1",
@@ -18,19 +17,17 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 model = AutoModelForCausalLM.from_pretrained(
     "cropinailab/aksara_v1",
-    torch_dtype=torch.bfloat16,
-    device_map="cuda",          # ensures GPU usage on Spaces
+    torch_dtype=torch.float32,   # CPU-safe dtype
 )
 
 model.eval()
 
-print("Model loaded successfully!")
+print("Model loaded successfully on CPU!")
 
 
 # ===============================
-# Generation Function (GPU)
+# Generation Function
 # ===============================
-@spaces.GPU
 def generate_agri_response(plant, disease):
     prompt = f"""
 You are an agricultural expert specializing in plant pathology, crop nutrition, and safe farm management.
@@ -135,29 +132,22 @@ Include:
 - Proper storage of harvested produce (if disease affects storage)
 """
 
-    # Tokenize
-    inputs = tokenizer(
-        prompt,
-        return_tensors="pt",
-    ).to(model.device)
-
-    # Generate
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=600,
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.15,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-
-    full_output = tokenizer.decode(
-        outputs[0],
-        skip_special_tokens=True,
-    )
-
-    # Remove echoed prompt safely
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    with torch.no_grad():  # important for CPU efficiency
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=600,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.15,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Remove echoed prompt
     if full_output.startswith(prompt):
         cleaned = full_output[len(prompt):].strip()
     else:
@@ -167,7 +157,7 @@ Include:
 
 
 # ===============================
-# Local Test (optional)
+# Test
 # ===============================
 if __name__ == "__main__":
     print(generate_agri_response("Potato", "Late Blight"))