ZeroGPU best practice: load models at module level (cuda), inference only inside @spaces.GPU

src/models/sql_generator.py

@@ -1,16 +1,16 @@
 """
 SQL Generator: text-to-SQL via the trained Qwen2.5-Coder-7B LoRA.
 
-Loads the merged 16-bit checkpoint from Hugging Face Hub by default.
-The same `DanielRegaladoCardoso/sql-generator-qwen25-coder-7b-lora` repo
-contains both the LoRA adapter and the merged model.
+Loads at module import time (root level), as required by ZeroGPU best
+practices. Inference happens inside @spaces.GPU in the orchestrator.
 """
 
 import logging
 import re
 from typing import Optional
 
-from src.models.base import BaseModel
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -21,11 +21,11 @@ SYSTEM_PROMPT = (
     "Return only the SQL."
 )
 
+DEFAULT_MODEL = "DanielRegaladoCardoso/sql-generator-qwen25-coder-7b-lora"
 
-class SQLGenerator(BaseModel):
-    """Text-to-SQL generator using the trained Qwen2.5-Coder-7B model."""
 
-    DEFAULT_MODEL = "DanielRegaladoCardoso/sql-generator-qwen25-coder-7b-lora"
+class SQLGenerator:
+    """Text-to-SQL generator. Model loaded at construction time onto CUDA."""
 
     def __init__(
         self,
@@ -33,33 +33,23 @@ class SQLGenerator(BaseModel):
         temperature: float = 0.0,
         max_new_tokens: int = 400,
     ) -> None:
-        super().__init__(model_name="sql-generator")
         self.hf_model = hf_model
         self.temperature = temperature
         self.max_new_tokens = max_new_tokens
 
-    def load(self) -> None:
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-        import torch
-
-        logger.info(f"Loading SQL generator: {self.hf_model}")
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        dtype = torch.bfloat16 if device == "cuda" else torch.float32
-
+        logger.info(f"Loading SQL generator at module level: {self.hf_model}")
         self.tokenizer = AutoTokenizer.from_pretrained(self.hf_model)
+        # On ZeroGPU, device_map='cuda' uses emulation mode at module load and
+        # real GPU inside @spaces.GPU calls.
         self.model = AutoModelForCausalLM.from_pretrained(
             self.hf_model,
-            torch_dtype=dtype,
-            device_map=device,
+            torch_dtype=torch.bfloat16,
+            device_map="cuda",
         )
         self.model.eval()
-        self.is_loaded = True
-        logger.info(f"SQL generator loaded on {device}")
-
-    def generate(self, question: str, schema: str) -> str:  # type: ignore[override]
-        self._validate_loaded()
-        import torch
+        logger.info("SQL generator ready")
 
+    def generate(self, question: str, schema: str) -> str:
         user_content = f"### Schema\n{schema}\n\n### Question\n{question}"
         messages = [
             {"role": "system", "content": SYSTEM_PROMPT},
@@ -81,18 +71,13 @@ class SQLGenerator(BaseModel):
         text = self.tokenizer.decode(
             out[0][input_ids.shape[1]:], skip_special_tokens=True
         )
-        sql = self._clean_sql(text)
-        logger.info(f"Generated SQL: {sql[:100]}...")
-        return sql
+        return self._clean_sql(text)
 
     @staticmethod
     def _clean_sql(text: str) -> str:
-        """Strip code fences, trailing prose, ensure a single SQL statement."""
         text = text.strip()
-        # Strip fences ```sql ... ```
         text = re.sub(r"^```(?:sql)?\s*", "", text, flags=re.IGNORECASE)
         text = re.sub(r"\s*```\s*$", "", text)
-        # Cut at the first ;  if followed by prose
         if ";" in text:
             stmt, _, _ = text.partition(";")
             text = stmt + ";"