arbor-ai 0.1.0__py3-none-any.whl

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/client/api.py

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

arbor/server/__init__.py

@@ -0,0 +1 @@
+

arbor/server/api/__init__.py

@@ -0,0 +1 @@
+

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/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/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/server/core/__init__.py

@@ -0,0 +1 @@
+

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/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/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/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/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/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()
+    }

arbor_ai-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,25 @@
+arbor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/cli.py,sha256=6fT5JjpXSwhpJSQNE4pnLOY04ryHPwJBAOet3hyho8k,383
+arbor/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/client/api.py,sha256=WFaNtwCNWXRAHHG1Jfyl7LvTP6jiEyQOLZn2Z8Yjt5k,40
+arbor/server/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+arbor/server/api/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+arbor/server/api/routes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/api/routes/files.py,sha256=QrPY9-886NXnXjGRlT-pl5kWbnwfogCrdmv6RufJpVg,466
+arbor/server/api/routes/jobs.py,sha256=ibL0tQA2Apqa91vycv3NPT0ydhkba4vnPoclw-bVKXs,510
+arbor/server/api/routes/training.py,sha256=43NOvh1Hubg3Ocfhu5E82Tp_kXOJL8H8oQXf-4H1yMU,981
+arbor/server/core/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+arbor/server/core/config.py,sha256=R67gNeUXz0RShvpr8XF3Lpn7-RMOfKf2xTIyqXvj4PI,215
+arbor/server/core/logging.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/main.py,sha256=I3chVYsoG56zE7Clf88lEuOPaDzJvKsOzivOWpsFDls,350
+arbor/server/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/services/dependencies.py,sha256=y3EoIkwScYc811jZ8p5m0kJT4ixRo7vguimBKKMuxAQ,458
+arbor/server/services/file_manager.py,sha256=uE7Mnbn9fC1H7sAMzJt1x9fVak10duh0OyAZeDgi3iY,2200
+arbor/server/services/job_manager.py,sha256=Zx3d0h31YH9bQ4yQr3FUXUGEHd-KUiTekZ0ndGOptrY,1893
+arbor/server/services/training_manager.py,sha256=SNumrzM1B-V1HBucUHmxnmc4rmhSCHVgPPc5nNPRC4Q,10682
+arbor/server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/utils/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor_ai-0.1.0.dist-info/METADATA,sha256=1KSvnCQMVYhdQ9ONUupWtO14MuEyirX2yD58i7xQfVc,1272
+arbor_ai-0.1.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+arbor_ai-0.1.0.dist-info/entry_points.txt,sha256=AaLg05CZSQeP2oGlCH_AnmZPz-zzLlVtpXToI4cM3kY,39
+arbor_ai-0.1.0.dist-info/RECORD,,

arbor_ai-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.1
+Root-Is-Purelib: true
+Tag: py3-none-any

arbor_ai-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+arbor=arbor.cli:cli
+