Load LoRA via PeftModel on top of standard base models (fixes r=16 vs r=8 mismatch)

src/models/svg_renderer.py

@@ -1,15 +1,16 @@
 """
-SVG Renderer: chart spec + data -> inline SVG.
-
-Model loaded at root module level (ZeroGPU best practice). If the model
-output isn't a valid SVG, falls back to themed Plotly.
+SVG Renderer: load the trained LoRA on top of DeepSeek Coder 1.3B base.
+Falls back to themed Plotly if the model output isn't a valid SVG.
 """
 
+import json
 import logging
+import re
 from typing import Any, Dict, List
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 
 from src.visualization.plotly_fallback import PlotlyRenderer
 from src.visualization.svg_theme import apply_theme, is_renderable_svg
@@ -23,43 +24,47 @@ SYSTEM_PROMPT = (
     "minimalist style. Return only the SVG, starting with <svg."
 )
 
-DEFAULT_MODEL = "DanielRegaladoCardoso/svg-renderer-deepseek-coder-1.3b-lora"
+BASE_MODEL = "deepseek-ai/deepseek-coder-1.3b-instruct"
+ADAPTER_REPO = "DanielRegaladoCardoso/svg-renderer-deepseek-coder-1.3b-lora"
 
 
 class SVGRenderer:
-    """Render a chart spec to inline SVG."""
-
-    def __init__(
-        self,
-        hf_model: str = DEFAULT_MODEL,
-        temperature: float = 0.2,
-        max_new_tokens: int = 1500,
-    ) -> None:
-        self.hf_model = hf_model
+
+    def __init__(self, temperature: float = 0.2, max_new_tokens: int = 1500) -> None:
         self.temperature = temperature
         self.max_new_tokens = max_new_tokens
         self._plotly = PlotlyRenderer()
+        self.model = None
+        self.tokenizer = None
 
-        logger.info(f"Loading SVG renderer at module level: {self.hf_model}")
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(self.hf_model)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.hf_model,
+            logger.info(f"Loading SVG base: {BASE_MODEL}")
+            self.tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+            base = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL,
                 torch_dtype=torch.bfloat16,
                 device_map="cuda",
+                trust_remote_code=True,
             )
+            # Try LoRA. If it fails (e.g., adapter has only model weights as one-piece file
+            # rather than a peft adapter), fall back to base model.
+            try:
+                self.model = PeftModel.from_pretrained(
+                    base,
+                    ADAPTER_REPO,
+                    torch_dtype=torch.bfloat16,
+                )
+                logger.info("SVG renderer ready (LoRA applied)")
+            except Exception as e:
+                logger.warning(f"LoRA load failed ({e}); using base model")
+                self.model = base
             self.model.eval()
-            logger.info("SVG renderer ready")
         except Exception as e:
-            logger.warning(f"SVG model load failed ({e}); will use Plotly fallback only")
+            logger.warning(f"SVG model load failed entirely ({e}); Plotly fallback only")
             self.model = None
             self.tokenizer = None
 
-    def generate(
-        self,
-        chart_spec: Dict[str, Any],
-        data: List[Dict[str, Any]],
-    ) -> str:
+    def generate(self, chart_spec: Dict[str, Any], data: List[Dict[str, Any]]) -> str:
         if self.model is not None and self.tokenizer is not None:
             try:
                 svg = self._generate_model(chart_spec, data)
@@ -72,11 +77,7 @@ class SVGRenderer:
         svg = self._plotly.render(chart_spec, data)
         return apply_theme(svg)
 
-    def _generate_model(
-        self, chart_spec: Dict[str, Any], data: List[Dict[str, Any]]
-    ) -> str:
-        import json
-
+    def _generate_model(self, chart_spec: Dict[str, Any], data: List[Dict[str, Any]]) -> str:
         sample = data[:50]
         user_content = (
             f"Chart spec: {json.dumps(chart_spec, default=str)}\n"
@@ -107,6 +108,5 @@ class SVGRenderer:
 
     @staticmethod
     def _extract_svg(text: str) -> str:
-        import re
         m = re.search(r"<svg[\s\S]*?</svg>", text, re.IGNORECASE)
         return m.group(0) if m else text.strip()