feat: download Qwen2.5-Coder-0.5B + BART at build, add few-shot prompts

Dockerfile

@@ -2,15 +2,28 @@ FROM python:3.10-slim
 
 WORKDIR /app
 
-COPY . .
+COPY requirements.txt .
 
 RUN pip install --no-cache-dir --upgrade pip && \
     pip install --no-cache-dir -r requirements.txt
 
-RUN mkdir -p /app/data/uploads /app/static/images
-
+# Pre-download models during build so runtime stays offline
 ENV TRANSFORMERS_CACHE=/app/.cache/huggingface/transformers
 ENV HF_HOME=/app/.cache/huggingface
+
+RUN python -c "\
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM; \
+AutoTokenizer.from_pretrained('ArchCoder/fine-tuned-bart-large'); \
+AutoModelForSeq2SeqLM.from_pretrained('ArchCoder/fine-tuned-bart-large'); \
+AutoTokenizer.from_pretrained('Qwen/Qwen2.5-Coder-0.5B-Instruct'); \
+AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-Coder-0.5B-Instruct'); \
+print('Models downloaded successfully')"
+
+COPY . .
+
+RUN mkdir -p /app/data/uploads /app/static/images
+
+# Lock to offline at runtime — all models are already cached
 ENV HF_HUB_OFFLINE=1
 ENV TRANSFORMERS_OFFLINE=1
 ENV HF_HUB_DISABLE_TELEMETRY=1

app.py

@@ -54,12 +54,12 @@ def index():
 def models():
     return jsonify({
         "models": [
-            {"id": "qwen",   "name": "Qwen2.5-1.5B",    "provider": "Local (optional path)", "free": True},
-            {"id": "bart",   "name": "BART (fine-tuned)", "provider": "Local (transformers)", "free": True},
-            {"id": "gemini", "name": "Gemini 2.0 Flash", "provider": "Google AI (API key)",  "free": False},
-            {"id": "grok",   "name": "Grok-3 Mini",      "provider": "xAI (API key)",       "free": False},
+            {"id": "qwen",   "name": "Qwen2.5-Coder-0.5B", "provider": "Local (transformers)", "free": True},
+            {"id": "bart",   "name": "BART (fine-tuned)",   "provider": "Local (transformers)", "free": True},
+            {"id": "gemini", "name": "Gemini 2.0 Flash",    "provider": "Google AI (API key)",  "free": False},
+            {"id": "grok",   "name": "Grok-3 Mini",         "provider": "xAI (API key)",       "free": False},
         ],
-        "default": "bart"
+        "default": "qwen"
     })
 
 
@@ -70,7 +70,7 @@ def plot():
     if not data or not data.get('query'):
         return jsonify({'error': 'Missing required field: query'}), 400
 
-    logging.info(f"Plot request: model={data.get('model','bart')} query={data.get('query')[:80]}")
+    logging.info(f"Plot request: model={data.get('model','qwen')} query={data.get('query')[:80]}")
     result = agent.process_request(data)
     logging.info(f"Plot completed in {time.time() - t0:.2f}s")
     return jsonify(result)

llm_agent.py

@@ -15,50 +15,69 @@ load_dotenv()
 
 logger = logging.getLogger(__name__)
 
+# ---------------------------------------------------------------------------
+# Model IDs (downloaded at Docker build, cached in HF_HOME)
+# ---------------------------------------------------------------------------
+QWEN_MODEL_ID = os.getenv("QWEN_MODEL_ID", "Qwen/Qwen2.5-Coder-0.5B-Instruct")
+BART_MODEL_ID = os.getenv("BART_MODEL_ID", "ArchCoder/fine-tuned-bart-large")
 
-def _model_dir(dirname: str) -> str:
-    return os.path.join(os.path.dirname(os.path.abspath(__file__)), dirname)
+# ---------------------------------------------------------------------------
+# Prompt templates with few-shot examples
+# ---------------------------------------------------------------------------
 
+_SYSTEM_PROMPT = """\
+You are a data visualization expert. Given the user request and dataset schema, \
+output ONLY a valid JSON object. No explanation, no markdown fences, no extra text.
 
-def _has_model_weights(model_dir: str) -> bool:
-    weight_files = (
-        "pytorch_model.bin",
-        "model.safetensors",
-        "tf_model.h5",
-        "flax_model.msgpack",
-    )
-    return os.path.isdir(model_dir) and any(
-        os.path.exists(os.path.join(model_dir, filename)) for filename in weight_files
-    )
+Required JSON keys:
+  "x"          : string  — exact column name for the x-axis
+  "y"          : array   — one or more exact column names for the y-axis
+  "chart_type" : string  — one of: line, bar, scatter, pie, histogram, box, area
+  "color"      : string or null — optional CSS color like "red", "#4f8cff"
 
-# ---------------------------------------------------------------------------
-# Prompt templates
-# ---------------------------------------------------------------------------
+Rules:
+- Use ONLY column names from the schema. Never invent names.
+- For pie charts: y must contain exactly one column.
+- For histogram/box: x may equal the first element of y.
+- Default to "line" if chart type is ambiguous.
+
+### Examples
+
+Example 1:
+Schema: Year (integer), Sales (float), Profit (float)
+User: "plot sales over the years with a red line"
+Output: {"x": "Year", "y": ["Sales"], "chart_type": "line", "color": "red"}
+
+Example 2:
+Schema: Month (string), Revenue (float), Expenses (float)
+User: "bar chart comparing revenue and expenses by month"
+Output: {"x": "Month", "y": ["Revenue", "Expenses"], "chart_type": "bar", "color": null}
+
+Example 3:
+Schema: Category (string), Count (integer)
+User: "pie chart of count by category"
+Output: {"x": "Category", "y": ["Count"], "chart_type": "pie", "color": null}
+
+Example 4:
+Schema: Date (string), Temperature (float), Humidity (float)
+User: "scatter plot of temperature vs humidity in blue"
+Output: {"x": "Temperature", "y": ["Humidity"], "chart_type": "scatter", "color": "blue"}
 
-_SYSTEM_PROMPT = (
-    "You are a data visualization expert. "
-    "Given the user request and the dataset schema provided, output ONLY a valid JSON "
-    "object — no explanation, no markdown fences, no extra text.\n\n"
-    "Required keys:\n"
-    '  "x"          : string  — exact column name for the x-axis\n'
-    '  "y"          : array   — one or more exact column names for the y-axis\n'
-    '  "chart_type" : string  — one of: line, bar, scatter, pie, histogram, box, area\n'
-    '  "color"      : string  — optional CSS color, e.g. "red", "#4f8cff"\n\n'
-    "Rules:\n"
-    "- Use only column names that appear in the schema. Never invent names.\n"
-    "- For pie: y must contain exactly one column.\n"
-    "- For histogram/box: x may equal the first element of y.\n"
-    "- Default to line if chart type is ambiguous."
-)
+Example 5:
+Schema: Year (integer), Sales (float), Employee expense (float), Marketing expense (float)
+User: "show me an area chart of sales and marketing expense over years"
+Output: {"x": "Year", "y": ["Sales", "Marketing expense"], "chart_type": "area", "color": null}
+"""
 
 
 def _user_message(query: str, columns: list, dtypes: dict, sample_rows: list) -> str:
     schema = "\n".join(f"  - {c} ({dtypes.get(c, 'unknown')})" for c in columns)
     samples = "".join(f"  {json.dumps(r)}\n" for r in sample_rows[:3])
     return (
-        f"Dataset columns:\n{schema}\n\n"
-        f"Sample rows (first 3):\n{samples}\n"
-        f"User request: {query}"
+        f"Schema:\n{schema}\n\n"
+        f"Sample rows:\n{samples}\n"
+        f"User: \"{query}\"\n"
+        f"Output:"
     )
 
 
@@ -198,22 +217,16 @@ class LLM_Agent:
         self._bart_model = None
         self._qwen_tokenizer = None
         self._qwen_model = None
-        self._bart_model_dir = os.getenv("BART_LOCAL_PATH", _model_dir("fine-tuned-bart-large"))
-        self._qwen_model_dir = os.getenv("QWEN_LOCAL_PATH", "")
 
     # -- model runners -------------------------------------------------------
 
     def _run_qwen(self, user_msg: str) -> str:
+        """Qwen2.5-Coder-0.5B-Instruct — fast structured-JSON generation."""
         if self._qwen_model is None:
             from transformers import AutoModelForCausalLM, AutoTokenizer
-            model_id = self._qwen_model_dir
-            if not model_id:
-                raise ValueError("Qwen local model is not configured in this Space")
-            if not _has_model_weights(model_id):
-                raise ValueError(f"Qwen model weights not found in {model_id}")
-            logger.info("Loading Qwen model (first request)...")
-            self._qwen_tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True)
-            self._qwen_model = AutoModelForCausalLM.from_pretrained(model_id, local_files_only=True)
+            logger.info(f"Loading Qwen model: {QWEN_MODEL_ID}")
+            self._qwen_tokenizer = AutoTokenizer.from_pretrained(QWEN_MODEL_ID)
+            self._qwen_model = AutoModelForCausalLM.from_pretrained(QWEN_MODEL_ID)
             logger.info("Qwen model loaded.")
         messages = [
             {"role": "system", "content": _SYSTEM_PROMPT},
@@ -260,14 +273,12 @@ class LLM_Agent:
         return resp.choices[0].message.content
 
     def _run_bart(self, query: str) -> str:
+        """ArchCoder/fine-tuned-bart-large — lightweight Seq2Seq fallback."""
         if self._bart_model is None:
             from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-            model_id = self._bart_model_dir
-            if not _has_model_weights(model_id):
-                raise ValueError(f"BART model weights not found in {model_id}")
-            logger.info("Loading BART model (first request)...")
-            self._bart_tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True)
-            self._bart_model = AutoModelForSeq2SeqLM.from_pretrained(model_id, local_files_only=True)
+            logger.info(f"Loading BART model: {BART_MODEL_ID}")
+            self._bart_tokenizer = AutoTokenizer.from_pretrained(BART_MODEL_ID)
+            self._bart_model = AutoModelForSeq2SeqLM.from_pretrained(BART_MODEL_ID)
             logger.info("BART model loaded.")
         inputs = self._bart_tokenizer(
             query, return_tensors="pt", max_length=512, truncation=True
@@ -281,7 +292,7 @@ class LLM_Agent:
         t0        = time.time()
         query     = data.get("query", "")
         data_path = data.get("file_path")
-        model     = data.get("model", "bart")
+        model     = data.get("model", "qwen")
 
         if data_path and os.path.exists(data_path):
             self.data_processor  = DataProcessor(data_path)
@@ -303,14 +314,15 @@ class LLM_Agent:
             user_msg = _user_message(query, columns, dtypes, sample_rows)
             if   model == "gemini": raw_text = self._run_gemini(user_msg)
             elif model == "grok":   raw_text = self._run_grok(user_msg)
+            elif model == "bart":   raw_text = self._run_bart(query)
             elif model == "qwen":
                 try:
                     raw_text = self._run_qwen(user_msg)
                 except Exception as qwen_exc:
-                    logger.warning(f"Qwen unavailable, falling back to BART: {qwen_exc}")
+                    logger.warning(f"Qwen failed, falling back to BART: {qwen_exc}")
                     raw_text = self._run_bart(query)
             else:
-                raw_text = self._run_bart(query)
+                raw_text = self._run_qwen(user_msg)
 
             logger.info(f"LLM [{model}] output: {raw_text}")
             parsed    = _parse_output(raw_text)