arbor-ai 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

arbor/cli.py

@@ -1,17 +1,101 @@
 import click
 import uvicorn
+
+from arbor.server.core.config import Settings
 from arbor.server.main import app
+from arbor.server.services.file_manager import FileManager
+from arbor.server.services.grpo_manager import GRPOManager
+from arbor.server.services.inference_manager import InferenceManager
+from arbor.server.services.job_manager import JobManager
+from arbor.server.services.training_manager import TrainingManager
+
 
 @click.group()
 def cli():
     pass
 
+
+def create_app(arbor_config_path: str):
+    """Create and configure the Arbor API application
+
+    Args:
+        storage_path (str): Path to store models and uploaded training files
+
+    Returns:
+        FastAPI: Configured FastAPI application
+    """
+    # Create new settings instance with overrides
+    settings = Settings.load_from_yaml(arbor_config_path)
+
+    # Initialize services with settings
+    file_manager = FileManager(settings=settings)
+    job_manager = JobManager(settings=settings)
+    training_manager = TrainingManager(settings=settings)
+    inference_manager = InferenceManager(settings=settings)
+    grpo_manager = GRPOManager(settings=settings)
+    # Inject settings into app state
+    app.state.settings = settings
+    app.state.file_manager = file_manager
+    app.state.job_manager = job_manager
+    app.state.training_manager = training_manager
+    app.state.inference_manager = inference_manager
+    app.state.grpo_manager = grpo_manager
+
+    return app
+
+
+def start_server(host="0.0.0.0", port=7453, storage_path="./storage", timeout=10):
+    """Start the Arbor API server with a single function call"""
+    import socket
+    import threading
+    import time
+    from contextlib import closing
+
+    def is_port_in_use(port):
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+            return sock.connect_ex(("localhost", port)) == 0
+
+    # First ensure the port is free
+    if is_port_in_use(port):
+        raise RuntimeError(f"Port {port} is already in use")
+
+    app = create_app(storage_path)
+    config = uvicorn.Config(app, host=host, port=port, log_level="info")
+    server = uvicorn.Server(config)
+
+    def run_server():
+        server.run()
+
+    thread = threading.Thread(target=run_server, daemon=True)
+    thread.start()
+
+    # Wait for server to start
+    start_time = time.time()
+    while not is_port_in_use(port):
+        if time.time() - start_time > timeout:
+            raise TimeoutError(f"Server failed to start within {timeout} seconds")
+        time.sleep(0.1)
+
+    # Give it a little extra time to fully initialize
+    time.sleep(0.5)
+
+    return server
+
+
+def stop_server(server):
+    """Stop the Arbor API server"""
+    server.should_exit = True
+
+
 @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):
+@click.option("--host", default="0.0.0.0", help="Host to bind to")
+@click.option("--port", default=7453, help="Port to bind to")
+@click.option("--arbor-config", required=True, help="Path to the Arbor config file")
+def serve(host, port, arbor_config):
     """Start the Arbor API server"""
+    app = create_app(arbor_config)
     uvicorn.run(app, host=host, port=port)
 
-if __name__ == '__main__':
-    cli()
+
+if __name__ == "__main__":
+    cli()

arbor_ai-0.1.5.dist-info/METADATA

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: arbor-ai
+Version: 0.1.5
+Summary: A framework for fine-tuning and managing language models
+Author-email: Noah Ziems <nziems2@nd.edu>
+Project-URL: Homepage, https://github.com/Ziems/arbor
+Project-URL: Issues, https://github.com/Ziems/arbor/issues
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: fastapi
+Requires-Dist: uvicorn
+Requires-Dist: click
+Requires-Dist: python-multipart
+Requires-Dist: pydantic-settings
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: trl
+Requires-Dist: peft
+Requires-Dist: ray>=2.9
+Requires-Dist: setuptools<77.0.0,>=76.0.0
+Requires-Dist: pyzmq>=26.4.0
+Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: sglang>=0.4.5.post3
+Requires-Dist: sglang-router
+Dynamic: license-file
+
+<p align="center">
+  <img src="https://github.com/user-attachments/assets/ed0dd782-65fa-48b5-a762-b343b183be09" alt="Description" width="400"/>
+</p>
+
+**A framework for optimizing DSPy programs with RL.**
+
+---
+
+## 🚀 Installation
+
+Install Arbor via pip:
+
+```bash
+pip install git+https://github.com/Ziems/arbor.git
+```
+
+---
+
+## ⚡ Quick Start
+
+### 1️⃣ Make an `arbor.yaml` File
+
+This is all dependent on your setup. Here is an example of one:
+```yaml
+inference:
+  gpu_ids: '0'
+
+training:
+  gpu_ids: '1, 2'
+```
+
+### 2️⃣ Start the Server
+
+**CLI:**
+
+```bash
+python -m arbor.cli serve --arbor-config arbor.yaml
+```
+
+### 3️⃣ Optimize a DSPy Program
+
+Follow the DSPy tutorials here to see usage examples:
+[DSPy RL Optimization Examples](https://dspy.ai/tutorials/rl_papillon/)
+
+---
+
+## 🙏 Acknowledgements
+
+Arbor builds on the shoulders of great work. We extend our thanks to:
+- **[Will Brown's Verifiers library](https://github.com/willccbb/verifiers)**
+- **[Hugging Face TRL library](https://github.com/huggingface/trl)**

arbor_ai-0.1.5.dist-info/RECORD

@@ -0,0 +1,8 @@
+arbor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/cli.py,sha256=3o9A03Kew9cM5ZvD_6xOTaquNIE_hTYMOeQH3hkuJbY,3110
+arbor_ai-0.1.5.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
+arbor_ai-0.1.5.dist-info/METADATA,sha256=Tney6uOytHDMIZg3iqKrn2lgtaF3NULjXo19XdG_2Dw,1823
+arbor_ai-0.1.5.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+arbor_ai-0.1.5.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
+arbor_ai-0.1.5.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
+arbor_ai-0.1.5.dist-info/RECORD,,

{arbor_ai-0.1.4.dist-info → arbor_ai-0.1.5.dist-info}/WHEEL

@@ -1,4 +1,5 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.1.1
+Generator: setuptools (80.3.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
+

arbor_ai-0.1.5.dist-info/entry_points.txt

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

arbor_ai-0.1.5.dist-info/top_level.txt

@@ -0,0 +1 @@
+arbor

arbor/client/api.py

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

arbor/server/__init__.py

@@ -1 +0,0 @@
-

arbor/server/api/__init__.py

@@ -1 +0,0 @@
-

arbor/server/api/models/schemas.py

@@ -1,19 +0,0 @@
-from pydantic import BaseModel
-
-class FileResponse(BaseModel):
-    id: str
-    object: str = "file"
-    bytes: int
-    created_at: int
-    filename: str
-    purpose: str
-
-class FineTuneRequest(BaseModel):
-    model: str
-    training_file: str  # id of uploaded jsonl file
-
-class JobStatusResponse(BaseModel):
-    id: str
-    status: str
-    details: str = ""
-    fine_tuned_model: str | None = None

arbor/server/api/routes/files.py

@@ -1,23 +0,0 @@
-from fastapi import APIRouter, UploadFile, File, Depends, HTTPException
-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
-from arbor.server.services.file_manager import FileValidationError
-
-router = APIRouter()
-
-@router.post("", response_model=FileResponse)
-async def upload_file(
-    file: UploadFile = File(...),
-    file_manager: FileManager = Depends(get_file_manager)
-):
-    if not file.filename.endswith('.jsonl'):
-        raise HTTPException(status_code=400, detail="Only .jsonl files are allowed")
-
-    try:
-        content = await file.read()
-        file_manager.validate_file_format(content)
-        await file.seek(0)  # Reset file pointer to beginning
-        return file_manager.save_uploaded_file(file)
-    except FileValidationError as e:
-        raise HTTPException(status_code=400, detail=f"Invalid file format: {str(e)}")

arbor/server/api/routes/jobs.py

@@ -1,14 +0,0 @@
-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)
-):
-    job = job_manager.get_job(job_id)
-    return JobStatusResponse(id=job_id, status=job.status.value, fine_tuned_model=job.fine_tuned_model)

arbor/server/api/routes/training.py

@@ -1,16 +0,0 @@
-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(id=job.id, status=job.status.value)

arbor/server/core/__init__.py

@@ -1 +0,0 @@
-

arbor/server/core/config.py

@@ -1,10 +0,0 @@
-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

@@ -1,10 +0,0 @@
-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

@@ -1,16 +0,0 @@
-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

@@ -1,128 +0,0 @@
-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 FileValidationError(Exception):
-    """Custom exception for file validation errors"""
-    pass
-
-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]
-
-  def validate_file_format(self, file_content: bytes) -> None:
-    """
-    Validates that the file content is properly formatted JSONL with expected structure.
-    Raises FileValidationError if validation fails.
-    """
-    if not file_content:
-      raise FileValidationError("File is empty")
-
-    try:
-      lines = file_content.decode('utf-8').strip().split('\n')
-      if not lines:
-        raise FileValidationError("File contains no valid data")
-
-      for line_num, line in enumerate(lines, 1):
-        try:
-          data = json.loads(line)
-
-          # Validate required structure
-          if not isinstance(data, dict):
-            raise FileValidationError(f"Line {line_num}: Each line must be a JSON object")
-
-          if "messages" not in data:
-            raise FileValidationError(f"Line {line_num}: Missing 'messages' field")
-
-          if not isinstance(data["messages"], list):
-            raise FileValidationError(f"Line {line_num}: 'messages' must be an array")
-
-          for msg in data["messages"]:
-            if not isinstance(msg, dict):
-              raise FileValidationError(f"Line {line_num}: Each message must be an object")
-            if "role" not in msg or "content" not in msg:
-              raise FileValidationError(f"Line {line_num}: Messages must have 'role' and 'content' fields")
-            if not isinstance(msg["role"], str) or not isinstance(msg["content"], str):
-              raise FileValidationError(f"Line {line_num}: Message 'role' and 'content' must be strings")
-
-        except json.JSONDecodeError:
-          raise FileValidationError(f"Invalid JSON on line {line_num}")
-
-    except UnicodeDecodeError:
-      raise FileValidationError("File must be valid UTF-8 encoded text")

arbor/server/services/job_manager.py

@@ -1,76 +0,0 @@
-import uuid
-from enum import Enum
-import logging
-from datetime import datetime
-
-# https://platform.openai.com/docs/api-reference/fine-tuning/object
-class JobStatus(Enum):
-    PENDING = "pending" # Not in OAI
-    VALIDATING_FILES = "validating_files"
-    QUEUED = "queued"
-    RUNNING = "running"
-    SUCCEEDED = "succeeded"
-    FAILED = "failed"
-    CANCELLED = "cancelled"
-
-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.fine_tuned_model = None
-    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(self, job_id: str):
-    if job_id not in self.jobs:
-      raise ValueError(f"Job {job_id} not found")
-    return self.jobs[job_id]
-
-  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

@@ -1,264 +0,0 @@
-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}" # 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
-            ).to(device)
-            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=request.model)
-
-            # 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.SUCCEEDED
-            job.fine_tuned_model = sft_config.output_dir
-        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.4.dist-info/METADATA

@@ -1,97 +0,0 @@
-Metadata-Version: 2.3
-Name: arbor-ai
-Version: 0.1.4
-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.9, <3.14
-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.9
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
-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 🌳
-
-A drop-in replacement for OpenAI's fine-tuning API that lets you fine-tune and manage open-source language models locally. Train and deploy custom models with the same API you already know.
-
-## Installation
-
-```bash
-pip install arbor-ai
-```
-
-## Quick Start
-
-1. Start the Arbor server:
-
-```bash
-arbor serve
-```
-
-2. The server will be available at `http://localhost:8000`.
-
-3. Upload your training data:
-
-```python
-import requests
-
-requests.post('http://127.0.0.1:8000/api/files', files={'file': open('your_file.jsonl', 'rb')})
-```
-
-4. Submit a fine-tuning job:
-
-```python
-requests.post('http://127.0.0.1:8000/api/fine-tune', json={'model': 'HuggingFaceTB/SmolLM2-135M-Instruct', 'training_file': 'Returned file ID from Step 3'})
-```
-
-5. Monitor the job status:
-
-```python
-requests.get('http://127.0.0.1:8000/api/jobs/{Returned job ID from Step 4}')
-```
-
-
-
-## Development Setup
-
-```bash
-poetry install
-```
-
-```bash
-poetry run arbor serve
-```
-
-```bash
-poetry run pytest
-```
-
-## Contributing
-
-Contributions are welcome! Please feel free to submit a Pull Request.
-
-## License
-
-This project is licensed under the MIT License - see the LICENSE file for details.
-
-## Support
-
-If you encounter any issues or have questions, please file an issue on the [GitHub repository](https://github.com/Ziems/arbor/issues).

arbor_ai-0.1.4.dist-info/RECORD

@@ -1,27 +0,0 @@
-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/models/schemas.py,sha256=19uDproKWhPQvVTit0hWuqmPb80zrELtCgnLybDuBKw,398
-arbor/server/api/routes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arbor/server/api/routes/files.py,sha256=U5QPC05VzqgDirB77lpy6BJLvg3zo1eGz7RUEk3HgRw,970
-arbor/server/api/routes/jobs.py,sha256=W2Y-rByaULxT0pEy3_YSNWO2CEKR5obyax-uR4ax_6Y,539
-arbor/server/api/routes/training.py,sha256=5M6OAtl9i8L-jBefmvPWvyf1M_x30-IlXzgleBg41Yc,977
-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=VUCn0cUtd-Q1BrUPtKStS1hGtV_OlymUyA0I8zeG9Po,4037
-arbor/server/services/job_manager.py,sha256=rZjuhwwbvL7yCJi653tv7z36iFFvp1w5J9j5DntSWKM,2073
-arbor/server/services/training_manager.py,sha256=BQsUsxOyRlgFDEFM77tyIahmm4NqcoOwxq8Tlmp66dY,10724
-arbor/server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arbor/server/utils/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arbor_ai-0.1.4.dist-info/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
-arbor_ai-0.1.4.dist-info/METADATA,sha256=977OGIuruJzS8wkFntELEoO7Ey5VzEhv88v1Pt81pa0,2451
-arbor_ai-0.1.4.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-arbor_ai-0.1.4.dist-info/entry_points.txt,sha256=AaLg05CZSQeP2oGlCH_AnmZPz-zzLlVtpXToI4cM3kY,39
-arbor_ai-0.1.4.dist-info/RECORD,,

arbor_ai-0.1.4.dist-info/entry_points.txt

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