Spaces:
Running
Running
| """ | |
| Minimal Training Script using Unsloth & TRL for the OpenEnv Supply Chain Environment. | |
| You can run this in a Google Colab notebook instance. | |
| """ | |
| # Install dependencies in Colab: | |
| # !pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" | |
| # !pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes | |
| # !pip install openenv-core | |
| import torch | |
| from unsloth import FastLanguageModel | |
| from trl import SFTTrainer, PPOConfig, PPOTrainer | |
| from transformers import TrainingArguments | |
| from openenv_wrapper import SupplyChainEnv | |
| # 1. Load Unsloth FastLanguageModel | |
| max_seq_length = 2048 | |
| model, tokenizer = FastLanguageModel.from_pretrained( | |
| model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit", | |
| max_seq_length = max_seq_length, | |
| dtype = None, | |
| load_in_4bit = True, | |
| ) | |
| # 2. Add LoRA adapters | |
| model = FastLanguageModel.get_peft_model( | |
| model, | |
| r = 16, | |
| target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", | |
| "gate_proj", "up_proj", "down_proj",], | |
| lora_alpha = 16, | |
| lora_dropout = 0, | |
| bias = "none", | |
| use_gradient_checkpointing = "unsloth", | |
| random_state = 3407, | |
| ) | |
| # 3. Setup OpenEnv Environment | |
| env = SupplyChainEnv(max_turns=12) | |
| # For actual Reinforcement Learning (TRL + PPO), you would use the PPOTrainer | |
| # wrapped around the environment rollout loops. | |
| # Here is a stub for the SFT/RL loop taking observations from the env: | |
| def run_rl_episode(env, model, tokenizer): | |
| obs, info = env.reset() | |
| done = False | |
| # Simple prompt template for the model | |
| system_prompt = "You are an AI Supply Chain Manager. Respond ONLY with valid JSON actions." | |
| while not done: | |
| prompt = f"<|system|>\n{system_prompt}\n<|user|>\nObservation:\n{obs}\n<|assistant|>\n" | |
| inputs = tokenizer(prompt, return_tensors="pt").to("cuda") | |
| outputs = model.generate(**inputs, max_new_tokens=100) | |
| # Decode the action generated by the LLM | |
| action_str = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) | |
| # Step the environment | |
| obs, reward, done, truncated, info = env.step(action_str) | |
| print(f"Action Taken: {action_str} | Reward: {reward}") | |
| # To train with SFT based on expert trajectories (if generated): | |
| # trainer = SFTTrainer( | |
| # model = model, | |
| # train_dataset = your_expert_dataset, | |
| # dataset_text_field = "text", | |
| # max_seq_length = max_seq_length, | |
| # args = TrainingArguments( | |
| # per_device_train_batch_size = 2, | |
| # gradient_accumulation_steps = 4, | |
| # warmup_steps = 5, | |
| # max_steps = 60, | |
| # learning_rate = 2e-4, | |
| # fp16 = not torch.cuda.is_bfloat16_supported(), | |
| # bf16 = torch.cuda.is_bfloat16_supported(), | |
| # logging_steps = 1, | |
| # optim = "adamw_8bit", | |
| # output_dir = "outputs", | |
| # ), | |
| # ) | |
| # trainer.train() | |
| if __name__ == "__main__": | |
| print("Executing a test rollout with the untrained LoRA model...") | |
| run_rl_episode(env, model, tokenizer) | |