refactor: split [evaluator] into evaluator (cloud) and evaluator-local (HuggingFace local) with lazy imports

Author: OdinHoang03

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "xfmr-zem"
-version = "0.3.3"
+version = "0.3.4"
 description = "Zem: Unified Data Pipeline Framework (ZenML + NeMo Curator + DataJuicer) for multi-domain processing"
 readme = "README.md"
 requires-python = ">=3.10,<3.13"
@@ -119,10 +119,16 @@ voice = [
 # ── LLM / Evaluation ──────────────────────────────────────────────────────────
 evaluator = [
     "opik>=1.10.9",
+    "litellm>=1.0.0",
+]
+
+# evaluator-local: chạy test/eval model local (HuggingFace)
+evaluator-local = [
+    "opik>=1.10.9",
+    "litellm>=1.0.0",
     "transformers>=4.40.0",
     "torch>=2.1.0",
     "accelerate>=0.25.0",
-    "litellm>=1.0.0",
 ]
 
 # ── Web UI ─────────────────────────────────────────────────────────────────────

src/xfmr_zem/servers/evaluator/factory/eval_engines/local.py

@@ -1,41 +1,53 @@
-import torch
 import json
 import re
 from typing import Any, Dict, List, Optional
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from opik import track
 from pydantic import BaseModel
 from loguru import logger
 import sys
+
 logger.remove()
 logger.add(sys.stderr, level="INFO")
-# Giả lập OpikBaseModel
 
 from opik.evaluation.models import OpikBaseModel
 
+
 class OpikHFModel(OpikBaseModel):
     def __init__(self, model_id: str, **kwargs):
+        # Lazy import: chỉ import khi thực sự cần (evaluator-local extra)
+        try:
+            import torch
+            from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline as hf_pipeline
+            self._torch = torch
+            self._AutoTokenizer = AutoTokenizer
+            self._AutoModelForCausalLM = AutoModelForCausalLM
+            self._hf_pipeline = hf_pipeline
+        except ImportError:
+            raise ImportError(
+                "Thiếu dependencies cho local model. "
+                "Hãy cài: pip install 'xfmr-zem[evaluator-local]'"
+            )
+
         super().__init__(model_name=model_id)
         self.model_id = model_id
         self.max_new_tokens = kwargs.get("max_new_tokens", 512)
-        # Temperature thấp để model tập trung vào logic chấm điểm, không sáng tạo lung tung
-        self.temperature = kwargs.get("temperature", 0.01) 
+        self.temperature = kwargs.get("temperature", 0.01)
         self.device = kwargs.get("device", "auto")
         self._load_model()
-    
+
     @track(name="load_hf_model")
     def _load_model(self):
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.tokenizer = self._AutoTokenizer.from_pretrained(self.model_id)
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
-            
-        self.model = AutoModelForCausalLM.from_pretrained(
+
+        self.model = self._AutoModelForCausalLM.from_pretrained(
             self.model_id,
-            torch_dtype=torch.float16,
+            torch_dtype=self._torch.float16,
             device_map=self.device
         )
-        
-        self.generator = pipeline(
+
+        self.generator = self._hf_pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
@@ -45,89 +57,72 @@ def _load_model(self):
     def _extract_json_string(self, text: str) -> str:
         """Cắt lấy phần JSON từ output"""
         text = text.strip()
-        # Regex tìm ```json { ... } ```
         match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL)
-        if match: return match.group(1)
-        
-        # Fallback: tìm { }
+        if match:
+            return match.group(1)
         start, end = text.find("{"), text.rfind("}")
-        if start != -1 and end != -1: return text[start : end + 1]
+        if start != -1 and end != -1:
+            return text[start: end + 1]
         return text
 
     @logger.catch(reraise=True)
     @track(name="hf_generate_string")
     def generate_string(self, input: str, response_format: Any = None, **kwargs: Any) -> str:
-        # 1. LOG INPUT
-        # Quan sát xem input thực tế nhận vào là gì (đã có schema hay chưa?)
         logger.info(f"\n--- [GENERATE START] ---\n")
 
-        # 2. GENERATION
-        # Không can thiệp sửa input nữa, chỉ setup tham số chạy
         params = {
             "max_new_tokens": kwargs.get("max_new_tokens", self.max_new_tokens),
-            "temperature": kwargs.get("temperature", 0.01 if response_format else self.temperature), # Temp thấp nếu cần JSON
+            "temperature": kwargs.get("temperature", 0.01 if response_format else self.temperature),
             "do_sample": True,
             "return_full_text": False
         }
 
         try:
-            # Gọi model sinh text
             response = self.generator(input, **params)
             raw_text = response[0]["generated_text"].strip()
-            
-            # Nếu không yêu cầu format đặc biệt, trả về luôn
+
             if not response_format:
                 return raw_text
 
-            # 3. CHECK FORMAT (VALIDATION)
             logger.info("--- [VALIDATING JSON] ---")
-            
-            # Bước A: Extract JSON từ text (lọc rác markdown)
             json_str = self._extract_json_string(raw_text)
-            
-            # Bước B: Parse & Validate
-            data = json.loads(json_str) # Thử parse JSON thuần
-            
-            # Nếu có Pydantic Model, validate chặt chẽ kiểu dữ liệu
+            data = json.loads(json_str)
+
             if isinstance(response_format, type) and issubclass(response_format, BaseModel):
                 logger.info(f"Validating against Pydantic Model: {response_format.__name__}")
                 validated_obj = response_format.model_validate(data)
-                
-                # Thành công!
                 final_json = validated_obj.model_dump_json()
                 logger.info("VALIDATION SUCCESS ✅")
                 return final_json
-            
-            # Nếu chỉ là dict schema thường
+
             logger.info("VALIDATION SUCCESS (Dict) ✅")
             return json.dumps(data, ensure_ascii=False)
 
         except json.JSONDecodeError as e:
             logger.error(f"❌ JSON PARSE ERROR: {e}\nBad String: {json_str}")
-            # Trả về lỗi dạng JSON để Opik ghi nhận thay vì crash chương trình
             return f'{{"error": "JSONDecodeError", "details": "{str(e)}", "raw_output": "{raw_text}"}}'
-            
+
         except Exception as e:
             logger.error(f"❌ GENERATION/VALIDATION ERROR: {e}")
             return f'{{"error": "RuntimeError", "details": "{str(e)}"}}'
 
     @track(name="hf_generate_provider_response")
     def generate_provider_response(self, messages: List[Dict[str, Any]], **kwargs: Any) -> Any:
-        # Chuyển messages thành prompt string cơ bản
         prompt = "\n".join([f"{m.get('role','').title()}: {m.get('content','')}" for m in messages])
         prompt += "\nAssistant:"
-        
+
         generated_text = self.generate_string(
-            prompt, 
-            response_format=kwargs.pop("response_format", None), 
+            prompt,
+            response_format=kwargs.pop("response_format", None),
             **kwargs
         )
-        
+
         return {
             "choices": [{"message": {"role": "assistant", "content": generated_text}}],
             "model": self.model_id
         }
 
+
 class OpikLocalFactory:
     @staticmethod
     def create_model(provider: str, model_id: str, **kwargs) -> Any:

src/xfmr_zem/servers/evaluator/factory/models.py

@@ -1,10 +1,19 @@
 from typing import Any, Optional
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 import opik
 
+
 class HuggingFaceLM:
     def __init__(self, model_id: str):
+        # Lazy import: chỉ import khi thực sự cần (evaluator-local extra)
+        try:
+            from transformers import AutoModelForCausalLM, AutoTokenizer
+            import torch
+        except ImportError:
+            raise ImportError(
+                "Thiếu dependencies cho local model. "
+                "Hãy cài: pip install 'xfmr-zem[evaluator-local]'"
+            )
+
         self.model_id = model_id
         print(f"Loading HuggingFace Model: {model_id}")
         self.model = AutoModelForCausalLM.from_pretrained(
@@ -13,7 +22,7 @@ def __init__(self, model_id: str):
             device_map="auto"
         )
         self.tokenizer = AutoTokenizer.from_pretrained(
-            model_id, 
+            model_id,
             torch_dtype="auto",
             device_map="auto"
         )
@@ -31,27 +40,28 @@ def generate(self, input_text: str, system_prompt: Optional[str] = None) -> str:
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": input_text}
         ]
-        
+
         text = self.tokenizer.apply_chat_template(
             messages,
             tokenize=False,
             add_generation_prompt=True
         )
-        
+
         model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
 
         generated_ids = self.model.generate(
             **model_inputs,
             max_new_tokens=512
         )
-        
+
         generated_ids = [
             output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
         ]
 
         response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return response
 
+
 class ModelFactory:
     @staticmethod
     def get_model(engine_type: str, model_id: str) -> Any: