arbor-ai 0.1.0__tar.gz

arbor_ai-0.1.0/PKG-INFO

@@ -0,0 +1,47 @@
+Metadata-Version: 2.3
+Name: arbor-ai
+Version: 0.1.0
+Summary: A framework for fine-tuning and managing language models
+License: MIT
+Keywords: machine learning,fine-tuning,language models
+Author: Noah Ziems
+Author-email: nziems2@nd.edu
+Requires-Python: >=3.13
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Requires-Dist: click
+Requires-Dist: fastapi
+Requires-Dist: peft (>=0.14.0,<0.15.0)
+Requires-Dist: pydantic-settings (>=2.8.1,<3.0.0)
+Requires-Dist: python-multipart (>=0.0.20,<0.0.21)
+Requires-Dist: torch (>=2.6.0,<3.0.0)
+Requires-Dist: transformers (>=4.49.0,<5.0.0)
+Requires-Dist: trl (>=0.15.2,<0.16.0)
+Requires-Dist: uvicorn
+Project-URL: Repository, https://github.com/arbor-ai/arbor
+Description-Content-Type: text/markdown
+
+# Arbor AI
+
+## Setup
+
+```bash
+poetry install
+```
+
+```bash
+poetry run arbor serve
+```
+
+## Uploading Data
+
+```bash
+curl -X POST "http://localhost:8000/api/files" -F "file=@training_data.jsonl"
+```
+

arbor_ai-0.1.0/README.md

@@ -0,0 +1,17 @@
+# Arbor AI
+
+## Setup
+
+```bash
+poetry install
+```
+
+```bash
+poetry run arbor serve
+```
+
+## Uploading Data
+
+```bash
+curl -X POST "http://localhost:8000/api/files" -F "file=@training_data.jsonl"
+```

arbor_ai-0.1.0/pyproject.toml

@@ -0,0 +1,52 @@
+[project]
+name = "arbor-ai"
+version = "0.1.0"
+description = "A framework for fine-tuning and managing language models"
+authors = [
+    {name = "Noah Ziems",email = "nziems2@nd.edu"}
+]
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = [
+    "fastapi",
+    "uvicorn",
+    "click",
+    "python-multipart (>=0.0.20,<0.0.21)",
+    "pydantic-settings (>=2.8.1,<3.0.0)",
+    "torch (>=2.6.0,<3.0.0)",
+    "transformers (>=4.49.0,<5.0.0)",
+    "trl (>=0.15.2,<0.16.0)",
+    "peft (>=0.14.0,<0.15.0)",
+]
+
+[project.scripts]
+arbor = "arbor.cli:cli"
+
+[tool.poetry]
+name = "arbor"
+version = "0.1.0"
+description = "A framework for fine-tuning and managing language models"
+authors = ["Noah Ziems <nziems2@nd.edu>"]
+readme = "README.md"
+packages = [{include = "arbor", from = "src"}]
+repository = "https://github.com/arbor-ai/arbor"
+keywords = ["machine learning", "fine-tuning", "language models"]
+license = "MIT"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+]
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.5"
+pytest-asyncio = "^0.25.3"
+httpx = "^0.28.1"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"

arbor_ai-0.1.0/src/arbor/cli.py

@@ -0,0 +1,17 @@
+import click
+import uvicorn
+from arbor.server.main import app
+
+@click.group()
+def cli():
+    pass
+
+@cli.command()
+@click.option('--host', default='0.0.0.0', help='Host to bind to')
+@click.option('--port', default=8000, help='Port to bind to')
+def serve(host, port):
+    """Start the Arbor API server"""
+    uvicorn.run(app, host=host, port=port)
+
+if __name__ == '__main__':
+    cli()

arbor_ai-0.1.0/src/arbor/client/api.py

@@ -0,0 +1,2 @@
+from typing import Optional, Dict, Any
+

arbor_ai-0.1.0/src/arbor/server/__init__.py

@@ -0,0 +1 @@
+

arbor_ai-0.1.0/src/arbor/server/api/__init__.py

@@ -0,0 +1 @@
+

arbor_ai-0.1.0/src/arbor/server/api/routes/files.py

@@ -0,0 +1,13 @@
+from fastapi import APIRouter, UploadFile, File, Depends
+from arbor.server.services.file_manager import FileManager
+from arbor.server.api.models.schemas import FileResponse
+from arbor.server.services.dependencies import get_file_manager
+
+router = APIRouter()
+
+@router.post("", response_model=FileResponse)
+def upload_file(
+    file: UploadFile = File(...),
+    file_manager: FileManager = Depends(get_file_manager)
+):
+    return file_manager.save_uploaded_file(file)

arbor_ai-0.1.0/src/arbor/server/api/routes/jobs.py

@@ -0,0 +1,14 @@
+from fastapi import APIRouter, Depends
+from arbor.server.services.job_manager import JobManager
+from arbor.server.services.dependencies import get_job_manager
+from arbor.server.api.models.schemas import JobStatusResponse
+
+router = APIRouter()
+
+@router.get("/{job_id}", response_model=JobStatusResponse)
+def get_job_status(
+    job_id: str,
+    job_manager: JobManager = Depends(get_job_manager)
+):
+    status = job_manager.get_job_status(job_id)
+    return JobStatusResponse(job_id=job_id, status=status.value)

arbor_ai-0.1.0/src/arbor/server/api/routes/training.py

@@ -0,0 +1,16 @@
+from fastapi import APIRouter, BackgroundTasks, Depends
+
+from arbor.server.api.models.schemas import FineTuneRequest, JobStatusResponse
+from arbor.server.services.job_manager import JobManager, JobStatus
+from arbor.server.services.file_manager import FileManager
+from arbor.server.services.training_manager import TrainingManager
+from arbor.server.services.dependencies import get_training_manager, get_job_manager, get_file_manager
+
+router = APIRouter()
+
+@router.post("", response_model=JobStatusResponse)
+def fine_tune(request: FineTuneRequest, background_tasks: BackgroundTasks, training_manager: TrainingManager = Depends(get_training_manager), job_manager: JobManager = Depends(get_job_manager), file_manager: FileManager = Depends(get_file_manager)):
+    job = job_manager.create_job()
+    background_tasks.add_task(training_manager.fine_tune, request, job, file_manager)
+    job.status = JobStatus.QUEUED
+    return JobStatusResponse(job_id=job.id, status=job.status.value)

arbor_ai-0.1.0/src/arbor/server/core/__init__.py

@@ -0,0 +1 @@
+

arbor_ai-0.1.0/src/arbor/server/core/config.py

@@ -0,0 +1,10 @@
+from pydantic_settings import BaseSettings
+
+class Settings(BaseSettings):
+    UPLOADS_DIR: str = "uploads"
+    MODEL_CACHE_DIR: str = "model_cache"
+
+    class Config:
+        env_file = ".env"
+
+settings = Settings()

arbor_ai-0.1.0/src/arbor/server/main.py

@@ -0,0 +1,10 @@
+from fastapi import FastAPI
+from arbor.server.api.routes import training, files, jobs
+from arbor.server.core.config import settings
+
+app = FastAPI(title="Arbor API")
+
+# Include routers
+app.include_router(training.router, prefix="/api/fine-tune")
+app.include_router(files.router, prefix="/api/files")
+app.include_router(jobs.router, prefix="/api/job")

arbor_ai-0.1.0/src/arbor/server/services/dependencies.py

@@ -0,0 +1,16 @@
+from functools import lru_cache
+from arbor.server.services.file_manager import FileManager
+from arbor.server.services.job_manager import JobManager
+from arbor.server.services.training_manager import TrainingManager
+
+@lru_cache()
+def get_file_manager() -> FileManager:
+    return FileManager()
+
+@lru_cache()
+def get_job_manager() -> JobManager:
+    return JobManager()
+
+@lru_cache()
+def get_training_manager() -> TrainingManager:
+    return TrainingManager()

arbor_ai-0.1.0/src/arbor/server/services/file_manager.py

@@ -0,0 +1,83 @@
+from pathlib import Path
+import json
+import os
+import shutil
+import time
+import uuid
+from fastapi import UploadFile
+from arbor.server.api.models.schemas import FileResponse
+
+class FileManager:
+  def __init__(self):
+    self.uploads_dir = Path("uploads")
+    self.uploads_dir.mkdir(exist_ok=True)
+    self.files = self.load_files_from_uploads()
+
+  def load_files_from_uploads(self):
+    files = {}
+
+    # Scan through all directories in uploads directory
+    for dir_path in self.uploads_dir.glob("*"):
+      if not dir_path.is_dir():
+        continue
+
+      # Check for metadata.json
+      metadata_path = dir_path / "metadata.json"
+      if not metadata_path.exists():
+        continue
+
+      # Load metadata
+      with open(metadata_path) as f:
+        metadata = json.load(f)
+
+      # Find the .jsonl file
+      jsonl_files = list(dir_path.glob("*.jsonl"))
+      if not jsonl_files:
+        continue
+
+      file_path = jsonl_files[0]
+      files[dir_path.name] = {
+        "path": str(file_path),
+        "purpose": metadata.get("purpose", "training"),
+        "bytes": file_path.stat().st_size,
+        "created_at": metadata.get("created_at", int(file_path.stat().st_mtime)),
+        "filename": metadata.get("filename", file_path.name)
+      }
+
+    return files
+
+  def save_uploaded_file(self, file: UploadFile) -> FileResponse:
+    file_id = str(uuid.uuid4())
+    dir_path = self.uploads_dir / file_id
+    dir_path.mkdir(exist_ok=True)
+
+    # Save the actual file
+    file_path = dir_path / f"data.jsonl"
+    with open(file_path, "wb") as f:
+      shutil.copyfileobj(file.file, f)
+
+    # Create metadata
+    metadata = {
+      "purpose": "training",
+      "created_at": int(time.time()),
+      "filename": file.filename
+    }
+
+    # Save metadata
+    with open(dir_path / "metadata.json", "w") as f:
+      json.dump(metadata, f)
+
+    file_data = {
+      "id": file_id,
+      "path": str(file_path),
+      "purpose": metadata["purpose"],
+      "bytes": file.size,
+      "created_at": metadata["created_at"],
+      "filename": metadata["filename"]
+    }
+
+    self.files[file_id] = file_data
+    return FileResponse(**file_data)
+
+  def get_file(self, file_id: str):
+    return self.files[file_id]

arbor_ai-0.1.0/src/arbor/server/services/job_manager.py

@@ -0,0 +1,72 @@
+import uuid
+from enum import Enum
+import logging
+from datetime import datetime
+
+class JobStatus(Enum):
+  PENDING = "pending"
+  QUEUED = "queued"
+  RUNNING = "running"
+  COMPLETED = "completed"
+  FAILED = "failed"
+
+class JobLogHandler(logging.Handler):
+  def __init__(self, job):
+    super().__init__()
+    self.job = job
+
+  def emit(self, record):
+    log_entry = {
+      'timestamp': datetime.fromtimestamp(record.created).isoformat(),
+      'level': record.levelname,
+      'message': record.getMessage()
+    }
+    self.job.logs.append(log_entry)
+
+class Job:
+  def __init__(self, id: str, status: JobStatus):
+    self.id = id
+    self.status = status
+    self.logs = []
+    self.logger = None
+    self.log_handler = None
+
+  def setup_logger(self, logger_name: str = None) -> logging.Logger:
+    """Sets up logging for the job with a dedicated handler."""
+    if logger_name is None:
+      logger_name = f"job_{self.id}"
+
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+
+    # Create and setup handler if not already exists
+    if self.log_handler is None:
+      handler = JobLogHandler(self)
+      formatter = logging.Formatter('%(message)s')
+      handler.setFormatter(formatter)
+      logger.addHandler(handler)
+      self.log_handler = handler
+
+    self.logger = logger
+    return logger
+
+  def cleanup_logger(self):
+    """Removes the job's logging handler."""
+    if self.logger and self.log_handler:
+      self.logger.removeHandler(self.log_handler)
+      self.log_handler = None
+      self.logger = None
+
+class JobManager:
+  def __init__(self):
+    self.jobs = {}
+
+  def get_job_status(self, job_id: str):
+    if job_id not in self.jobs:
+      raise ValueError(f"Job {job_id} not found")
+    return self.jobs[job_id].status
+
+  def create_job(self):
+    job = Job(id=str(uuid.uuid4()), status=JobStatus.PENDING)
+    self.jobs[job.id] = job
+    return job

arbor_ai-0.1.0/src/arbor/server/services/training_manager.py

@@ -0,0 +1,263 @@
+from arbor.server.api.models.schemas import FineTuneRequest
+from arbor.server.services.job_manager import Job, JobStatus
+from arbor.server.services.file_manager import FileManager
+
+class TrainingManager:
+    def __init__(self):
+        pass
+
+    def find_train_args(self, request: FineTuneRequest, file_manager: FileManager):
+        file = file_manager.get_file(request.training_file)
+        if file is None:
+            raise ValueError(f"Training file {request.training_file} not found")
+
+        data_path = file["path"]
+        output_dir = f"models/{request.model_name}" # TODO: This should be updated to be unique in some way
+
+
+        default_train_kwargs = {
+            "device": None,
+            "use_peft": False,
+            "num_train_epochs": 5,
+            "per_device_train_batch_size": 1,
+            "gradient_accumulation_steps": 8,
+            "learning_rate": 1e-5,
+            "max_seq_length": None,
+            "packing": True,
+            "bf16": True,
+            "output_dir": output_dir,
+            "train_data_path": data_path,
+        }
+        train_kwargs = {'packing': False}
+        train_kwargs={**default_train_kwargs, **(train_kwargs or {})}
+        output_dir = train_kwargs["output_dir"]  # user might have changed the output_dir
+
+        return train_kwargs
+
+
+    def fine_tune(self, request: FineTuneRequest, job: Job, file_manager: FileManager):
+        # Get logger for this job
+        logger = job.setup_logger("training")
+
+        job.status = JobStatus.RUNNING
+        logger.info("Starting fine-tuning job")
+
+        try:
+            train_kwargs = self.find_train_args(request, file_manager)
+
+            import torch
+            from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
+            from trl import SFTConfig, SFTTrainer, setup_chat_format
+
+            device = train_kwargs.get("device", None)
+            if device is None:
+                device = (
+                    "cuda"
+                    if torch.cuda.is_available()
+                    else "mps" if torch.backends.mps.is_available() else "cpu"
+                )
+            logger.info(f"Using device: {device}")
+
+            model = AutoModelForCausalLM.from_pretrained(
+                pretrained_model_name_or_path=request.model_name
+            ).to(device)
+            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=request.model_name)
+
+            # Set up the chat format; generally only for non-chat model variants, hence the try-except.
+            try:
+                model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)
+            except Exception:
+                pass
+
+            if tokenizer.pad_token_id is None:
+                logger.info("Adding pad token to tokenizer")
+                tokenizer.add_special_tokens({"pad_token": "[!#PAD#!]"})
+
+            logger.info("Creating dataset")
+            if "max_seq_length" not in train_kwargs or train_kwargs["max_seq_length"] is None:
+                train_kwargs["max_seq_length"] = 4096
+                logger.info(f"The 'train_kwargs' parameter didn't include a 'max_seq_length', defaulting to {train_kwargs['max_seq_length']}")
+
+
+            hf_dataset = dataset_from_file(train_kwargs["train_data_path"])
+            def tokenize_function(example):
+                return encode_sft_example(example, tokenizer, train_kwargs["max_seq_length"])
+            tokenized_dataset = hf_dataset.map(tokenize_function, batched=False)
+            tokenized_dataset.set_format(type="torch")
+            tokenized_dataset = tokenized_dataset.filter(lambda example: (example["labels"] != -100).any())
+
+            USE_PEFT = train_kwargs.get("use_peft", False)
+            peft_config = None
+
+            if USE_PEFT:
+                from peft import LoraConfig
+
+                rank_dimension = 32
+                lora_alpha = 64
+                lora_dropout = 0.05
+
+                peft_config = LoraConfig(
+                    r=rank_dimension,
+                    lora_alpha=lora_alpha,
+                    lora_dropout=lora_dropout,
+                    bias="none",
+                    target_modules="all-linear",
+                    task_type="CAUSAL_LM",
+                )
+
+            sft_config = SFTConfig(
+                output_dir=train_kwargs["output_dir"],
+                num_train_epochs=train_kwargs["num_train_epochs"],
+                per_device_train_batch_size=train_kwargs["per_device_train_batch_size"],
+                gradient_accumulation_steps=train_kwargs["gradient_accumulation_steps"],
+                learning_rate=train_kwargs["learning_rate"],
+                max_grad_norm=2.0,  # note that the current SFTConfig default is 1.0
+                logging_steps=20,
+                warmup_ratio=0.03,
+                lr_scheduler_type="constant",
+                save_steps=10_000,
+                bf16=train_kwargs["bf16"],
+                max_seq_length=train_kwargs["max_seq_length"],
+                packing=train_kwargs["packing"],
+                dataset_kwargs={  # We need to pass dataset_kwargs because we are processing the dataset ourselves
+                    "add_special_tokens": False,  # Special tokens handled by template
+                    "append_concat_token": False,  # No additional separator needed
+                },
+            )
+
+            logger.info("Starting training")
+            trainer = SFTTrainer(
+                model=model,
+                args=sft_config,
+                train_dataset=tokenized_dataset,
+                peft_config=peft_config,
+
+            )
+
+            # Train!
+            trainer.train()
+
+            # Save the model!
+            trainer.save_model()
+
+            MERGE = True
+            if USE_PEFT and MERGE:
+                from peft import AutoPeftModelForCausalLM
+
+                # Load PEFT model on CPU
+                model_ = AutoPeftModelForCausalLM.from_pretrained(
+                    pretrained_model_name_or_path=sft_config.output_dir,
+                    torch_dtype=torch.float16,
+                    low_cpu_mem_usage=True,
+                )
+
+                merged_model = model_.merge_and_unload()
+                merged_model.save_pretrained(
+                    sft_config.output_dir, safe_serialization=True, max_shard_size="5GB"
+                )
+
+            # Clean up!
+            import gc
+
+            del model
+            del tokenizer
+            del trainer
+            gc.collect()
+            torch.cuda.empty_cache()
+
+            logger.info("Training completed successfully")
+            job.status = JobStatus.COMPLETED
+        except Exception as e:
+            logger.error(f"Training failed: {str(e)}")
+            job.status = JobStatus.FAILED
+            raise
+        finally:
+            job.cleanup_logger()
+
+        return sft_config.output_dir
+
+def dataset_from_file(data_path):
+    """
+    Creates a HuggingFace Dataset from a JSONL file.
+    """
+    from datasets import load_dataset
+
+    dataset = load_dataset("json", data_files=data_path, split="train")
+    return dataset
+
+
+def encode_sft_example(example, tokenizer, max_seq_length):
+    """
+    This function encodes a single example into a format that can be used for sft training.
+    Here, we assume each example has a 'messages' field. Each message in it is a dict with 'role' and 'content' fields.
+    We use the `apply_chat_template` function from the tokenizer to tokenize the messages and prepare the input and label tensors.
+
+    Code obtained from the allenai/open-instruct repository: https://github.com/allenai/open-instruct/blob/4365dea3d1a6111e8b2712af06b22a4512a0df88/open_instruct/finetune.py
+    """
+    import torch
+
+    messages = example["messages"]
+    if len(messages) == 0:
+        raise ValueError("messages field is empty.")
+    input_ids = tokenizer.apply_chat_template(
+        conversation=messages,
+        tokenize=True,
+        return_tensors="pt",
+        padding=False,
+        truncation=True,
+        max_length=max_seq_length,
+        add_generation_prompt=False,
+    )
+    labels = input_ids.clone()
+    # mask the non-assistant part for avoiding loss
+    for message_idx, message in enumerate(messages):
+        if message["role"] != "assistant":
+            # we calculate the start index of this non-assistant message
+            if message_idx == 0:
+                message_start_idx = 0
+            else:
+                message_start_idx = tokenizer.apply_chat_template(
+                    conversation=messages[:message_idx],  # here marks the end of the previous messages
+                    tokenize=True,
+                    return_tensors="pt",
+                    padding=False,
+                    truncation=True,
+                    max_length=max_seq_length,
+                    add_generation_prompt=False,
+                ).shape[1]
+            # next, we calculate the end index of this non-assistant message
+            if message_idx < len(messages) - 1 and messages[message_idx + 1]["role"] == "assistant":
+                # for intermediate messages that follow with an assistant message, we need to
+                # set `add_generation_prompt=True` to avoid the assistant generation prefix being included in the loss
+                # (e.g., `<|assistant|>`)
+                message_end_idx = tokenizer.apply_chat_template(
+                    conversation=messages[: message_idx + 1],
+                    tokenize=True,
+                    return_tensors="pt",
+                    padding=False,
+                    truncation=True,
+                    max_length=max_seq_length,
+                    add_generation_prompt=True,
+                ).shape[1]
+            else:
+                # for the last message or the message that doesn't follow with an assistant message,
+                # we don't need to add the assistant generation prefix
+                message_end_idx = tokenizer.apply_chat_template(
+                    conversation=messages[: message_idx + 1],
+                    tokenize=True,
+                    return_tensors="pt",
+                    padding=False,
+                    truncation=True,
+                    max_length=max_seq_length,
+                    add_generation_prompt=False,
+                ).shape[1]
+            # set the label to -100 for the non-assistant part
+            labels[:, message_start_idx:message_end_idx] = -100
+            if max_seq_length and message_end_idx >= max_seq_length:
+                break
+    attention_mask = torch.ones_like(input_ids)
+    return {
+        "input_ids": input_ids.flatten(),
+        "labels": labels.flatten(),
+        "attention_mask": attention_mask.flatten()
+    }