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

arbor/server/api/models/schemas.py

@@ -9,10 +9,11 @@ class FileResponse(BaseModel):
     purpose: str
 
 class FineTuneRequest(BaseModel):
-    model_name: str
+    model: str
     training_file: str  # id of uploaded jsonl file
 
 class JobStatusResponse(BaseModel):
-    job_id: str
+    id: str
     status: str
-    details: str = ""
+    details: str = ""
+    fine_tuned_model: str | None = None

arbor/server/api/routes/files.py

@@ -1,13 +1,23 @@
-from fastapi import APIRouter, UploadFile, File, Depends
+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)
-def upload_file(
+async def upload_file(
     file: UploadFile = File(...),
     file_manager: FileManager = Depends(get_file_manager)
 ):
-    return file_manager.save_uploaded_file(file)
+    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

@@ -10,5 +10,5 @@ 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)
+    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

@@ -13,4 +13,4 @@ def fine_tune(request: FineTuneRequest, background_tasks: BackgroundTasks, train
     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)
+    return JobStatusResponse(id=job.id, status=job.status.value)

arbor/server/services/file_manager.py

@@ -7,6 +7,10 @@ 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")
@@ -80,4 +84,45 @@ class FileManager:
     return FileResponse(**file_data)
 
   def get_file(self, file_id: str):
-    return self.files[file_id]
+    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

@@ -3,12 +3,15 @@ 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"
-  QUEUED = "queued"
-  RUNNING = "running"
-  COMPLETED = "completed"
-  FAILED = "failed"
+    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):
@@ -27,6 +30,7 @@ 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
@@ -61,10 +65,10 @@ class JobManager:
   def __init__(self):
     self.jobs = {}
 
-  def get_job_status(self, job_id: str):
+  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].status
+    return self.jobs[job_id]
 
   def create_job(self):
     job = Job(id=str(uuid.uuid4()), status=JobStatus.PENDING)

arbor/server/services/training_manager.py

@@ -12,7 +12,7 @@ class TrainingManager:
             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
+        output_dir = f"models/{request.model}" # TODO: This should be updated to be unique in some way
 
 
         default_train_kwargs = {
@@ -59,9 +59,9 @@ class TrainingManager:
             logger.info(f"Using device: {device}")
 
             model = AutoModelForCausalLM.from_pretrained(
-                pretrained_model_name_or_path=request.model_name
+                pretrained_model_name_or_path=request.model
             ).to(device)
-            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=request.model_name)
+            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:
@@ -166,7 +166,8 @@ class TrainingManager:
             torch.cuda.empty_cache()
 
             logger.info("Training completed successfully")
-            job.status = JobStatus.COMPLETED
+            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

arbor_ai-0.1.4.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Noah Ziems
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

arbor_ai-0.1.4.dist-info/METADATA

@@ -0,0 +1,97 @@
+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.3.dist-info → arbor_ai-0.1.4.dist-info}/RECORD

@@ -4,23 +4,24 @@ 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=-_NOnUuEbiO8uLDwEpByYV6NAMasOmFUJXxG0eXA_D0,367
+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=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/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=uE7Mnbn9fC1H7sAMzJt1x9fVak10duh0OyAZeDgi3iY,2200
-arbor/server/services/job_manager.py,sha256=Zx3d0h31YH9bQ4yQr3FUXUGEHd-KUiTekZ0ndGOptrY,1893
-arbor/server/services/training_manager.py,sha256=SNumrzM1B-V1HBucUHmxnmc4rmhSCHVgPPc5nNPRC4Q,10682
+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.3.dist-info/METADATA,sha256=PJbOddt69fyZXJggzpaSUb5XfUt0ouPrQQFIAaeOasE,1272
-arbor_ai-0.1.3.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-arbor_ai-0.1.3.dist-info/entry_points.txt,sha256=AaLg05CZSQeP2oGlCH_AnmZPz-zzLlVtpXToI4cM3kY,39
-arbor_ai-0.1.3.dist-info/RECORD,,
+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.3.dist-info/METADATA

@@ -1,47 +0,0 @@
-Metadata-Version: 2.3
-Name: arbor-ai
-Version: 0.1.3
-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"
-```
-