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

arbor/server/services/training_manager.py

@@ -1,19 +1,40 @@
+import os
+import random
+import string
+import time
+from datetime import datetime
+from pathlib import Path
+
 from arbor.server.api.models.schemas import FineTuneRequest
-from arbor.server.services.job_manager import Job, JobStatus
+from arbor.server.core.config import Settings
 from arbor.server.services.file_manager import FileManager
+from arbor.server.services.job_manager import Job, JobEvent, JobStatus
 
-class TrainingManager:
-    def __init__(self):
-        pass
 
-    def find_train_args(self, request: FineTuneRequest, file_manager: FileManager):
+class TrainingManager:
+    def __init__(self, settings: Settings):
+        self.settings = settings
+
+    def make_output_dir(self, request: FineTuneRequest):
+        model_name = request.model.split("/")[-1].lower()
+        suffix = (
+            request.suffix
+            if request.suffix is not None
+            else "".join(random.choices(string.ascii_letters + string.digits, k=6))
+        )
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        name = f"ft:{model_name}:{suffix}:{timestamp}"
+        return name, str(Path(self.settings.STORAGE_PATH).resolve() / "models" / name)
+
+    def find_train_args_sft(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
+        file_manager.validate_file_format_sft(data_path)
 
+        name, output_dir = self.make_output_dir(request)
 
         default_train_kwargs = {
             "device": None,
@@ -28,25 +49,225 @@ class TrainingManager:
             "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
+        train_kwargs = {"packing": False}
+        train_kwargs = {**default_train_kwargs, **(train_kwargs or {})}
 
         return train_kwargs
 
+    def find_train_args_dpo(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"]
+        file_manager.validate_file_format_dpo(data_path)
+
+        name, output_dir = self.make_output_dir(request)
+
+        default_train_kwargs = {
+            "device": "cuda:2",
+            "use_peft": False,
+            "num_train_epochs": 5,
+            "per_device_train_batch_size": 1,
+            "gradient_accumulation_steps": 8,
+            "learning_rate": 1e-5,
+            "packing": True,
+            "bf16": True,
+            "output_dir": output_dir,
+            "train_data_path": data_path,
+            "prompt_length": 1024,
+            "max_seq_length": 1512,
+            "use_peft": False,
+        }
+
+        # https://www.philschmid.de/dpo-align-llms-in-2024-with-trl#3-align-llm-with-trl-and-the-dpotrainer
+
+        train_kwargs = request.model_dump(exclude_unset=True)
+        train_kwargs = {**default_train_kwargs, **(train_kwargs or {})}
+
+        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")
+        job.add_event(
+            JobEvent(level="info", message="Starting fine-tuning job", data={})
+        )
+
+        fine_tune_type = request.method["type"]
+        if fine_tune_type == "dpo":
+            self.dpo_fine_tune(request, job, file_manager)
+        else:
+            self.sft_fine_tune(request, job, file_manager)
+
+    def dpo_fine_tune(
+        self, request: FineTuneRequest, job: Job, file_manager: FileManager
+    ):
+        try:
+
+            job.status = JobStatus.RUNNING
+            job.add_event(
+                JobEvent(level="info", message="Starting fine-tuning job", data={})
+            )
+
+            train_kwargs = self.find_train_args_dpo(request, file_manager)
+            import torch
+            from transformers import (
+                AutoModelForCausalLM,
+                AutoTokenizer,
+                TrainingArguments,
+            )
+            from trl import DPOConfig, DPOTrainer, 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"
+                )
+
+            job.add_event(
+                JobEvent(level="info", message=f"Using device: {device}", data={})
+            )
+
+            model = AutoModelForCausalLM.from_pretrained(
+                pretrained_model_name_or_path=request.model, device_map="auto"
+            )
+            tokenizer = AutoTokenizer.from_pretrained(
+                pretrained_model_name_or_path=request.model
+            )
+
+            try:
+                model, tokenizer = setup_chat_format(model=model, tokenizer=tokenizer)
+            except Exception:
+                pass
+
+            if tokenizer.pad_token_id is None:
+                job.add_event(
+                    JobEvent(
+                        level="info", message="Adding pad token to tokenizer", data={}
+                    )
+                )
+                tokenizer.add_special_tokens({"pad_token": "[!#PAD#!]"})
+
+            hf_dataset = dataset_from_file(train_kwargs["train_data_path"])
+            train_dataset = hf_dataset
+
+            use_peft = train_kwargs.get("use_peft", False)
+            peft_config = None
+
+            if use_peft:
+                from peft import LoraConfig
+
+                peft_config = LoraConfig(
+                    lora_alpha=128,
+                    lora_dropout=0.05,
+                    r=256,
+                    bias="none",
+                    target_modules="all-linear",
+                    task_type="CAUSAL_LM",
+                )
+
+            dpo_args = DPOConfig(
+                output_dir=train_kwargs["output_dir"],
+                num_train_epochs=train_kwargs["num_train_epochs"],
+            )
+
+            trainer = DPOTrainer(
+                model=model,
+                args=dpo_args,
+                train_dataset=train_dataset,
+                processing_class=tokenizer,
+                peft_config=peft_config,
+            )
+
+            trainer.train()
+
+            if job.status == JobStatus.PENDING_PAUSE:
+                trainer.accelerator.save_state(output_dir=dpo_args.output_dir)
+                current_step = trainer.state.global_step
+                job.status = JobStatus.PAUSED
+                job.add_event(
+                    JobEvent(
+                        level="info",
+                        message="Training paused",
+                        data={"step": current_step},
+                    )
+                )
+
+            while job.status == JobStatus.PAUSED:
+                time.sleep(1)  # Sleep to avoid busy waiting
+                if job.status == JobStatus.PENDING_RESUME:
+                    job.status = JobStatus.RUNNING
+                    job.add_event(JobEvent(level="info", message="Resuming training"))
+                    trainer.accelerator.load_state(input_dir=dpo_args.output_dir)
+                    trainer.train(resume_from_checkpoint=True)
+
+            if job.status == JobStatus.PENDING_CANCEL:
+                job.status = JobStatus.CANCELLED
+                job.add_event(JobEvent(level="info", message="Training cancelled"))
+
+                _cleanup(model, tokenizer, trainer)
+                raise Exception(
+                    "Training cancelled"
+                )  # not sure if this should be raised or just return None
+
+            job.add_event(
+                JobEvent(level="info", message="Training completed successfully")
+            )
+
+            job.add_event(JobEvent(level="info", message="Saving model", data={}))
+            # Save the model!
+            trainer.save_model()
+            job.add_event(
+                JobEvent(
+                    level="info",
+                    message="Model saved",
+                    data={"location": dpo_args.output_dir},
+                )
+            )
+
+            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=dpo_args.output_dir,
+                    torch_dtype=torch.float16,
+                    low_cpu_mem_usage=True,
+                )
+
+                merged_model = model_.merge_and_unload()
+                merged_model.save_pretrained(
+                    dpo_args.output_dir, safe_serialization=True, max_shard_size="5GB"
+                )
+
+            _cleanup(model, tokenizer, trainer)
+
+            job.status = JobStatus.SUCCEEDED
+            job.fine_tuned_model = dpo_args.output_dir
+        except Exception as e:
+            job.add_event(
+                JobEvent(level="error", message=f"Training failed: {str(e)}", data={})
+            )
+            job.status = JobStatus.FAILED
+            raise
+        finally:
+            pass
+
+        return dpo_args.output_dir
+
+    def sft_fine_tune(
+        self, request: FineTuneRequest, job: Job, file_manager: FileManager
+    ):
 
         try:
-            train_kwargs = self.find_train_args(request, file_manager)
+            train_kwargs = self.find_train_args_sft(request, file_manager)
 
             import torch
-            from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
+            from transformers import AutoModelForCausalLM, AutoTokenizer
             from trl import SFTConfig, SFTTrainer, setup_chat_format
 
             device = train_kwargs.get("device", None)
@@ -56,12 +277,16 @@ class TrainingManager:
                     if torch.cuda.is_available()
                     else "mps" if torch.backends.mps.is_available() else "cpu"
                 )
-            logger.info(f"Using device: {device}")
+            job.add_event(
+                JobEvent(level="info", message=f"Using device: {device}", data={})
+            )
 
             model = AutoModelForCausalLM.from_pretrained(
                 pretrained_model_name_or_path=request.model
             ).to(device)
-            tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=request.model)
+            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:
@@ -70,21 +295,40 @@ class TrainingManager:
                 pass
 
             if tokenizer.pad_token_id is None:
-                logger.info("Adding pad token to tokenizer")
+                job.add_event(
+                    JobEvent(
+                        level="info", message="Adding pad token to tokenizer", data={}
+                    )
+                )
                 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:
+            job.add_event(JobEvent(level="info", message="Creating dataset", data={}))
+            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']}")
-
+                job.add_event(
+                    JobEvent(
+                        level="info",
+                        message=f"The 'train_kwargs' parameter didn't include a 'max_seq_length', defaulting to {train_kwargs['max_seq_length']}",
+                        data={},
+                    )
+                )
 
+            job.add_event(JobEvent(level="info", message="Tokenizing dataset", data={}))
             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"])
+                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())
+            tokenized_dataset = tokenized_dataset.filter(
+                lambda example: (example["labels"] != -100).any()
+            )
 
             USE_PEFT = train_kwargs.get("use_peft", False)
             peft_config = None
@@ -125,20 +369,61 @@ class TrainingManager:
                 },
             )
 
-            logger.info("Starting training")
+            job.add_event(JobEvent(level="info", message="Starting training"))
             trainer = SFTTrainer(
                 model=model,
                 args=sft_config,
                 train_dataset=tokenized_dataset,
                 peft_config=peft_config,
-
             )
 
             # Train!
             trainer.train()
 
+            if job.status == JobStatus.PENDING_PAUSE:
+                trainer.accelerator.save_state(output_dir=sft_config.output_dir)
+                current_step = trainer.state.global_step
+                job.status = JobStatus.PAUSED
+                job.add_event(
+                    JobEvent(
+                        level="info",
+                        message="Training paused",
+                        data={"step": current_step},
+                    )
+                )
+
+            while job.status == JobStatus.PAUSED:
+                time.sleep(1)  # Sleep to avoid busy waiting
+                if job.status == JobStatus.PENDING_RESUME:
+                    job.status = JobStatus.RUNNING
+                    job.add_event(JobEvent(level="info", message="Resuming training"))
+                    trainer.accelerator.load_state(input_dir=sft_config.output_dir)
+                    trainer.train(resume_from_checkpoint=True)
+
+            if job.status == JobStatus.PENDING_CANCEL:
+                job.status = JobStatus.CANCELLED
+                job.add_event(JobEvent(level="info", message="Training cancelled"))
+
+                _cleanup(model, tokenizer, trainer)
+
+                raise Exception(
+                    "Training cancelled"
+                )  # not sure if this should be raised or just return None
+
+            job.add_event(
+                JobEvent(level="info", message="Training completed successfully")
+            )
+
+            job.add_event(JobEvent(level="info", message="Saving model", data={}))
             # Save the model!
             trainer.save_model()
+            job.add_event(
+                JobEvent(
+                    level="info",
+                    message="Model saved",
+                    data={"location": sft_config.output_dir},
+                )
+            )
 
             MERGE = True
             if USE_PEFT and MERGE:
@@ -156,27 +441,22 @@ class TrainingManager:
                     sft_config.output_dir, safe_serialization=True, max_shard_size="5GB"
                 )
 
-            # Clean up!
-            import gc
+            _cleanup(model, tokenizer, trainer)
 
-            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.add_event(
+                JobEvent(level="error", message=f"Training failed: {str(e)}", data={})
+            )
             job.status = JobStatus.FAILED
             raise
         finally:
-            job.cleanup_logger()
+            pass
 
         return sft_config.output_dir
 
+
 def dataset_from_file(data_path):
     """
     Creates a HuggingFace Dataset from a JSONL file.
@@ -218,7 +498,9 @@ def encode_sft_example(example, tokenizer, max_seq_length):
                 message_start_idx = 0
             else:
                 message_start_idx = tokenizer.apply_chat_template(
-                    conversation=messages[:message_idx],  # here marks the end of the previous messages
+                    conversation=messages[
+                        :message_idx
+                    ],  # here marks the end of the previous messages
                     tokenize=True,
                     return_tensors="pt",
                     padding=False,
@@ -227,7 +509,10 @@ def encode_sft_example(example, tokenizer, 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":
+            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|>`)
@@ -260,5 +545,17 @@ def encode_sft_example(example, tokenizer, max_seq_length):
     return {
         "input_ids": input_ids.flatten(),
         "labels": labels.flatten(),
-        "attention_mask": attention_mask.flatten()
-    }
+        "attention_mask": attention_mask.flatten(),
+    }
+
+
+def _cleanup(model, tokenizer, trainer):
+    import gc
+
+    import torch
+
+    del model
+    del tokenizer
+    del trainer
+    gc.collect()
+    torch.cuda.empty_cache()

arbor_ai-0.1.6.dist-info/METADATA

@@ -0,0 +1,78 @@
+Metadata-Version: 2.4
+Name: arbor-ai
+Version: 0.1.6
+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.6.dist-info/RECORD

@@ -0,0 +1,34 @@
+arbor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/cli.py,sha256=3o9A03Kew9cM5ZvD_6xOTaquNIE_hTYMOeQH3hkuJbY,3110
+arbor/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/client/api.py,sha256=86bgHuGM_AvI1Uhic_QaCnpF4VFqXie9ZzxmbTXUPpQ,19
+arbor/server/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+arbor/server/main.py,sha256=tY4Vlaaj4oq1FTGYOkbFMGF0quLEeR-VBaKaXhQ5mEE,382
+arbor/server/api/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+arbor/server/api/models/schemas.py,sha256=s_G8sSb05FjkKEqpKpLlqaEd8NysJddHibRHhcnrKIk,5594
+arbor/server/api/routes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/api/routes/files.py,sha256=DQC_ogH5zlzhHZSAA4Cj5wzK07XBIBVs2Po91W9rcDY,1835
+arbor/server/api/routes/grpo.py,sha256=VuEvSOwwrHegn9qM-1nbHFmmUnnC_BMwnIHsfIdiJyI,1877
+arbor/server/api/routes/inference.py,sha256=xlP-FMpOJAiiPZkE470l9mCR0ujLki8RrcO9hmTQD-k,1662
+arbor/server/api/routes/jobs.py,sha256=BNdaSYUBJX6xSd6Pj6qx1DQJiZ5EKVxxbXDbEkfkCpw,3634
+arbor/server/core/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+arbor/server/core/config.py,sha256=Mx77S3ByIMvHmPDikQLcczhzA5so3Vrw_U4QefOiHOU,1257
+arbor/server/core/logging.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/services/dependencies.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/services/file_manager.py,sha256=Z9z4A4EzvPauid_DBfpim401DDtuJy_TbX4twTWDJWI,12119
+arbor/server/services/grpo_manager.py,sha256=T-f1TrNSi_kmxPOcpaDphS8Xf3UMUbricocc6fuaKIM,12077
+arbor/server/services/inference_manager.py,sha256=qR9xPiYs4Is24vgeF72w7Hbe8j_PGEbl-qewcvUV-dA,9731
+arbor/server/services/job_manager.py,sha256=m_d4UPwN_82f7t7K443DaFpFoyv7JZSZKml8tawt1Bk,2186
+arbor/server/services/training_manager.py,sha256=oQdhpfxdgp_lCTb_lxhvjupdLrcg6HL3TEbct_q9F6I,21065
+arbor/server/services/comms/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/services/comms/comms.py,sha256=Dg08D2Fm5TAEiGyr0Qcr0uocabQpFD_sBVhxIkj9D2M,7424
+arbor/server/services/scripts/grpo_training.py,sha256=V36pCMZDJj2DdzquxScOddi9zP8EVPGWN3HGiftFfrY,21082
+arbor/server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor/server/utils/helpers.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+arbor_ai-0.1.6.dist-info/licenses/LICENSE,sha256=5vFGrbOFeXXM83JV9o16w7ohH4WLeu3-57GocJSz8ow,1067
+arbor_ai-0.1.6.dist-info/METADATA,sha256=ot1XsjoFawGbNBmaAaxv29lE_uNh8TTdl8OANkTPTS8,1823
+arbor_ai-0.1.6.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+arbor_ai-0.1.6.dist-info/entry_points.txt,sha256=PGBX-MfNwfIl8UPFgsX3gjtXLqSogRhOktKMpZUysD0,40
+arbor_ai-0.1.6.dist-info/top_level.txt,sha256=jzWdp3BRYqvZDMFsPajrcftvvlluzVDErkD8IMRfhYs,6
+arbor_ai-0.1.6.dist-info/RECORD,,

{arbor_ai-0.1.4.dist-info → arbor_ai-0.1.6.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.6.dist-info/entry_points.txt

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

arbor_ai-0.1.6.dist-info/top_level.txt

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

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)