Merge pull request #279 from DashAISoftware/develop

Removed llama-cpp from requirements

Author: cristian-tamblay

DashAI/back/api/api_v1/endpoints/pipelines.py

@@ -22,6 +22,7 @@
 from DashAI.back.dependencies.database.models import Dataset, Pipeline
 from DashAI.back.dependencies.registry.component_registry import ComponentRegistry
 from DashAI.back.exploration.base_explorer import BaseExplorer
+from DashAI.back.pipeline.validator.nodes_definitions import NODES
 from DashAI.back.pipeline.validator.pipeline_validator import PipelineValidator
 from DashAI.back.pipeline.validator.validator import VALIDATOR_MAP
 
@@ -61,37 +62,22 @@ async def get_pipelines(
 
 @router.get("/nodes")
 async def get_nodes() -> List[Dict[str, Any]]:
-    """Retrieve pipeline node definitions.
-
-    Returns
-    -------
-    List[Dict[str, Any]]
-        A list of node definitions from the nodes.json file.
-
-    Raises
-    ------
-    HTTPException
-        500: Failed to load node definitions.
-    """
+    """Retrieve pipeline node definitions."""
     try:
-        json_path = Path(__file__).resolve().parents[3] / "pipeline" / "nodes.json"
-        with open(json_path, "r") as f:
-            nodes = json.load(f)
-
-        type_to_name = {node["type"]: node["name"] for node in nodes}
-        for node in nodes:
-            successors = node.get("successors", [])
-            node["next"] = [type_to_name.get(s, s) for s in successors]
+        type_to_name = {node.type: node.name for node in NODES}
+        nodes_with_next = []
+        for node in NODES:
+            node_dict = node.model_dump()
+            node_dict["next"] = [type_to_name.get(s, s) for s in node.successors]
+            nodes_with_next.append(node_dict)
+        return nodes_with_next
 
     except Exception as e:
-        logger.exception(e)
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail="Failed to load node definitions",
         ) from e
 
-    return nodes
-
 
 @router.get("/predict_summary")
 @inject

DashAI/back/dataloaders/classes/dashai_dataset.py

@@ -297,8 +297,16 @@ def save_dataset(dataset: DashAIDataset, path: Union[str, os.PathLike]) -> None:
         writer.close()
 
     metadata_filepath = os.path.join(path, "splits.json")
-    with open(metadata_filepath, "w") as f:
-        json.dump(dataset.splits, f, indent=2, sort_keys=True, ensure_ascii=False)
+    # Update splits with dataset shape and column names
+    metadata = dataset.splits
+    metadata.update(
+        {
+            "total_rows": dataset.shape[0],
+            "column_names": dataset.column_names,
+        }
+    )
+    with open(metadata_filepath, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2, sort_keys=True, ensure_ascii=False)
 
 
 @beartype
@@ -767,24 +775,16 @@ def get_dataset_info(dataset_path: str) -> object:
     else:
         splits_data = {"split_indices": {}}
 
-    data_filepath = os.path.join(dataset_path, "data.arrow")
-    with pa.OSFile(data_filepath, "rb") as source:
-        reader = ipc.open_file(source)
-        schema = reader.schema
-        column_names = schema.names
-
-        total_rows = 0
-        for i in range(reader.num_record_batches):
-            total_rows += reader.get_batch(i).num_rows
-
     splits = splits_data.get("split_indices", {})
     train_indices = splits.get("train", [])
     test_indices = splits.get("test", [])
     val_indices = splits.get("validation", [])
+    total_rows = splits_data.get("total_rows", 0)
+    column_names = splits_data.get("column_names", [])
 
     return {
         "total_rows": total_rows,
-        "total_columns": len(schema),
+        "total_columns": len(column_names),
         "column_names": column_names,
         "train_size": len(train_indices),
         "test_size": len(test_indices),

DashAI/back/initial_components.py

@@ -76,7 +76,6 @@
     DecisionTreeClassifier,
     DistilBertTransformer,
     DummyClassifier,
-    GemmaModel,
     GradientBoostingR,
     HistGradientBoostingClassifier,
     KNeighborsClassifier,
@@ -145,7 +144,6 @@ def get_initial_components():
         GradientBoostingR,
         HistGradientBoostingClassifier,
         KNeighborsClassifier,
-        GemmaModel,
         QwenModel,
         StableDiffusionV2Model,
         StableDiffusionV3Model,

DashAI/back/models/__init__.py

@@ -2,7 +2,6 @@
 from DashAI.back.models.base_generative_model import BaseGenerativeModel
 from DashAI.back.models.base_model import BaseModel
 from DashAI.back.models.hugging_face.distilbert_transformer import DistilBertTransformer
-from DashAI.back.models.hugging_face.gemma_model import GemmaModel
 from DashAI.back.models.hugging_face.opus_mt_en_es_transformer import (
     OpusMtEnESTransformer,
 )

DashAI/back/models/hugging_face/gemma_model.py

@@ -2,19 +2,20 @@
 import os
 from pathlib import Path
 
-import llama_cpp
 from packaging.version import Version
 
 logger = logging.getLogger(__name__)
 
+try:
+    import llama_cpp
+except ImportError:
+    llama_cpp = None
+
 
 def is_gpu_available_for_llama_cpp() -> bool:
-    """
-    Utility method to check if GPU offloading is supported for Llama models.
+    if llama_cpp is None:
+        return False
 
-    Returns:
-        bool: True if GPU offloading is supported, False otherwise or if fails.
-    """
     try:
         if Version(llama_cpp.__version__) > Version("0.3.0"):
             return __is_gpu_available_for_llama_cpp_v03()
@@ -23,7 +24,7 @@ def is_gpu_available_for_llama_cpp() -> bool:
 
     except Exception as e:
         logger.warning(
-            "Error checking GPU availability for llama_cpp. Will use CPU only. \n"
+            "Error checking GPU availability for llama_cpp. Will use CPU only.\n"
             f"Details: {e}"
         )
         return False

DashAI/back/models/hugging_face/llama_utils.py

@@ -1,6 +1,9 @@
 from typing import List
 
-from llama_cpp import Llama
+try:
+    from llama_cpp import Llama
+except ImportError:
+    Llama = None
 
 from DashAI.back.core.schema_fields import (
     BaseSchema,
@@ -14,7 +17,7 @@
     TextToTextGenerationTaskModel,
 )
 
-if is_gpu_available_for_llama_cpp():
+if Llama is not None and is_gpu_available_for_llama_cpp():
     DEVICE_ENUM = ["gpu", "cpu"]
     DEVICE_PLACEHOLDER = "gpu"
 else:
@@ -83,6 +86,11 @@ class QwenModel(TextToTextGenerationTaskModel):
     SCHEMA = QwenSchema
 
     def __init__(self, **kwargs):
+        if Llama is None:
+            raise RuntimeError(
+                "llama-cpp-python is not installed. Please install it to use QwenModel."
+            )
+
         kwargs = self.validate_and_transform(kwargs)
         self.model_name = kwargs.get("model_name", "Qwen/Qwen2.5-1.5B-Instruct-GGUF")
         self.max_tokens = kwargs.pop("max_tokens", 100)
@@ -107,6 +115,4 @@ def generate(self, prompt: list[dict[str, str]]) -> List[str]:
             temperature=self.temperature,
             frequency_penalty=self.frequency_penalty,
         )
-
-        generated_text = output["choices"][0]["message"]["content"]
-        return [generated_text]
+        return [output["choices"][0]["message"]["content"]]