Merge pull request #382 from MervinPraison/develop

Add Vision Model Training Support and Update Unsloth Dependency
2026-11-08 01:08:39 -10:00 · 2025-02-09 00:31:40 +00:00 · 2025-02-09 00:31:20 +00:00 · 2025-02-07 13:14:32 +00:00 · 2025-02-07 13:14:16 +00:00 · 2025-02-06 07:48:00 +00:00
11 changed files with 528 additions and 51 deletions
@@ -1,6 +1,6 @@
 FROM python:3.11-slim
 WORKDIR /app
 COPY . .
-RUN pip install flask praisonai==2.0.66 gunicorn markdown
+RUN pip install flask praisonai==2.0.75 gunicorn markdown
 EXPOSE 8080
 CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]
@@ -110,7 +110,7 @@ Loads environment variables from .env file or system and sets them.</p>
            file.write(&#34;FROM python:3.11-slim\n&#34;)
            file.write(&#34;WORKDIR /app\n&#34;)
            file.write(&#34;COPY . .\n&#34;)
-            file.write(&#34;RUN pip install flask praisonai==2.0.66 gunicorn markdown\n&#34;)
+            file.write(&#34;RUN pip install flask praisonai==2.0.75 gunicorn markdown\n&#34;)
            file.write(&#34;EXPOSE 8080\n&#34;)
            file.write(&#39;CMD [&#34;gunicorn&#34;, &#34;-b&#34;, &#34;0.0.0.0:8080&#34;, &#34;api:app&#34;]\n&#39;)
            
@@ -20,6 +20,20 @@ icon: "dumbbell"
 export HF_TOKEN=xxxxxxxxxxxxxx
 ```

+## Initilise praisonai train
+
+```bash
+praisonai train init
+```
+
+## Required
+
+1. Huggingface token
+2. Base model to train on (e.g. unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit)
+3. Dataset to train on (e.g. yahma/alpaca-cleaned)
+4. Huggingface model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
+5. Ollama model name to upload to (e.g. mervinpraison/llama3.1-instruct) (Optional)
+
 ## To upload to ollama.com (Linux)

 ```bash
@@ -108,4 +122,21 @@ quantization_method:

 ```bash
 praisonai train
+```
+
+## wandb
+
+```bash
+wandb login
+```
+
+<Note>
+Get the key from [here](https://wandb.ai/site/login)
+</Note>
+
+```bash
+export PRAISON_WANDB=True
+export WANDB_LOG_MODEL=checkpoint
+export WANDB_PROJECT=praisonai-test
+export PRAISON_WANDB_RUN_NAME=praisonai-train  
 ```
@@ -3,7 +3,7 @@ class Praisonai < Formula
  
    desc "AI tools for various AI applications"
    homepage "https://github.com/MervinPraison/PraisonAI"
-    url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/2.0.66.tar.gz"
+    url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/2.0.75.tar.gz"
    sha256 "1828fb9227d10f991522c3f24f061943a254b667196b40b1a3e4a54a8d30ce32"  # Replace with actual SHA256 checksum
    license "MIT"
  
@@ -189,8 +189,7 @@ class PraisonAI:
            package_root = os.path.dirname(os.path.abspath(__file__))
            config_yaml_destination = os.path.join(os.getcwd(), 'config.yaml')

-            # Create config.yaml only if it doesn't exist or --model or --dataset is provided
-            if not os.path.exists(config_yaml_destination) or args.model or args.dataset:
+            if not os.path.exists(config_yaml_destination):
                config = generate_config(
                    model_name=args.model,
                    hf_model_name=args.hf,
@@ -201,6 +200,20 @@ class PraisonAI:
                )
                with open('config.yaml', 'w') as f:
                    yaml.dump(config, f, default_flow_style=False, indent=2)
+            elif args.model or args.hf or args.ollama or (args.dataset and args.dataset != "yahma/alpaca-cleaned"):
+                config = generate_config(
+                    model_name=args.model,
+                    hf_model_name=args.hf,
+                    ollama_model_name=args.ollama,
+                    dataset=[{
+                        "name": args.dataset
+                    }]
+                )
+                with open('config.yaml', 'w') as f:
+                    yaml.dump(config, f, default_flow_style=False, indent=2)
+            else:
+                with open(config_yaml_destination, 'r') as f:
+                    config = yaml.safe_load(f)

            # Overwrite huggingface_save and ollama_save if --hf or --ollama are provided
            if args.hf:
@@ -227,7 +240,18 @@ class PraisonAI:
                print("All packages installed.")

            train_args = sys.argv[2:]  # Get all arguments after 'train'
-            train_script_path = os.path.join(package_root, 'train.py')
+            
+            # Check if this is a vision model - handle all case variations
+            model_name = config.get("model_name", "").lower()
+            is_vision_model = any(x in model_name for x in ["-vl-", "-vl", "vl-", "-vision-", "-vision", "vision-", "visionmodel"])
+            
+            # Choose appropriate training script
+            if is_vision_model:
+                train_script_path = os.path.join(package_root, 'train_vision.py')
+                print("Using vision training script for VL model...")
+            else:
+                train_script_path = os.path.join(package_root, 'train.py')
+                print("Using standard training script...")

            # Set environment variables
            env = os.environ.copy()
@@ -56,7 +56,7 @@ class CloudDeployer:
            file.write("FROM python:3.11-slim\n")
            file.write("WORKDIR /app\n")
            file.write("COPY . .\n")
-            file.write("RUN pip install flask praisonai==2.0.66 gunicorn markdown\n")
+            file.write("RUN pip install flask praisonai==2.0.75 gunicorn markdown\n")
            file.write("EXPOSE 8080\n")
            file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
            
@@ -96,10 +96,10 @@ PIP_FULL_PATH=$(conda run -n "$ENV_NAME" which pip)
 # Install other packages using pip
 $PIP_FULL_PATH install --upgrade pip 
 $PIP_FULL_PATH install "xformers==0.0.26.post1"
-$PIP_FULL_PATH install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git@038e6d4c8d40207a87297ab3aaf787c19b1006d1"
+$PIP_FULL_PATH install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git@53a773e4fbc53a1d96c7ba107e5fe75dab07027b"
 $PIP_FULL_PATH install --no-deps "trl<0.9.0" peft accelerate bitsandbytes 
 $PIP_FULL_PATH install unsloth_zoo
 $PIP_FULL_PATH install cut_cross_entropy
-$PIP_FULL_PATH install sentencepiece protobuf datasets huggingface_hub hf_transfer
+$PIP_FULL_PATH install sentencepiece protobuf datasets huggingface_hub hf_transfer wandb

 echo "Setup completed successfully!"
@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 """
-This script finetunes a model using Unsloth’s fast training framework.
+This script finetunes a model using Unsloth's fast training framework.
 It supports both ShareGPT and Alpaca‑style datasets by converting raw conversation
 data into plain-text prompts using a chat template, then pre‑tokenizing the prompts.
 Extra debug logging is added to help trace the root cause of errors.
@@ -27,7 +27,7 @@ from functools import partial
 #####################################
 def formatting_prompts_func(examples, tokenizer):
    """
-    Converts each example’s conversation into a single plain-text prompt.
+    Converts each example's conversation into a single plain-text prompt.
    If the example has a "conversations" field, process it as ShareGPT-style.
    Otherwise, assume Alpaca-style data with "instruction", "input", and "output" fields.
    """
@@ -216,23 +216,29 @@ class TrainModel:
        raw_dataset = self.load_datasets()
        tokenized_dataset = self.tokenize_dataset(raw_dataset)
        print("DEBUG: Dataset tokenization complete.")
-        training_args = TrainingArguments(
-            per_device_train_batch_size=2,
-            gradient_accumulation_steps=4,
-            warmup_steps=5,
-            max_steps=60,
-            learning_rate=2e-4,
-            fp16=not is_bfloat16_supported(),
-            bf16=is_bfloat16_supported(),
-            logging_steps=1,
-            optim="adamw_8bit",
-            weight_decay=0.01,
-            lr_scheduler_type="linear",
-            seed=3407,
-            output_dir="outputs",
-            report_to="none",
-            remove_unused_columns=False,
-        )
+        # Build the training arguments parameters dynamically
+        ta_params = {
+            "per_device_train_batch_size": self.config.get("per_device_train_batch_size", 2),
+            "gradient_accumulation_steps": self.config.get("gradient_accumulation_steps", 2),
+            "warmup_steps": self.config.get("warmup_steps", 50),
+            "max_steps": self.config.get("max_steps", 2800),
+            "learning_rate": self.config.get("learning_rate", 2e-4),
+            "fp16": self.config.get("fp16", not is_bfloat16_supported()),
+            "bf16": self.config.get("bf16", is_bfloat16_supported()),
+            "logging_steps": self.config.get("logging_steps", 15),
+            "optim": self.config.get("optim", "adamw_8bit"),
+            "weight_decay": self.config.get("weight_decay", 0.01),
+            "lr_scheduler_type": self.config.get("lr_scheduler_type", "linear"),
+            "seed": self.config.get("seed", 3407),
+            "output_dir": self.config.get("output_dir", "outputs"),
+            "report_to": "none" if not os.getenv("PRAISON_WANDB") else "wandb",
+            "remove_unused_columns": self.config.get("remove_unused_columns", False)
+        }
+        if os.getenv("PRAISON_WANDB"):
+            ta_params["save_steps"] = self.config.get("save_steps", 100)
+            ta_params["run_name"] = os.getenv("PRAISON_WANDB_RUN_NAME", "praisonai-train")
+
+        training_args = TrainingArguments(**ta_params)
        # Since the dataset is pre-tokenized, we supply a dummy dataset_text_field.
        trainer = SFTTrainer(
            model=self.model,
@@ -312,26 +318,208 @@ class TrainModel:

    def prepare_modelfile_content(self):
        output_model = self.config["hf_model_name"]
-        gguf_path = f"{output_model}/unsloth.Q4_K_M.gguf"
-        if not os.path.exists(gguf_path):
-            self.model, self.hf_tokenizer = self.load_model()
-            self.save_model_gguf()
-        return f"""FROM {output_model}/unsloth.Q4_K_M.gguf
-
-TEMPLATE \"\"\"Below are some instructions that describe some tasks. Write responses that appropriately complete each request.{{{{ if .Prompt }}}}
-
-### Instruction:
-{{{{ .Prompt }}}}
-
-{{{{ end }}}}### Response:
-{{{{ .Response }}}}\"\"\"
-
-PARAMETER stop ""
-PARAMETER stop ""
-PARAMETER stop ""
-PARAMETER stop ""
-PARAMETER stop "<|reserved_special_token_"
-"""
+        model_name = self.config["model_name"].lower()
+        # Mapping from model name keywords to their default TEMPLATE and stop tokens (and optional SYSTEM/num_ctx)
+        mapping = {
+            "llama": {
+                "template": """<|start_header_id|>system<|end_header_id|>
+    Cutting Knowledge Date: December 2023
+    {{ if .System }}{{ .System }}
+    {{- end }}
+    {{- if .Tools }}When you receive a tool call response, use the output to format an answer to the orginal user question.
+    You are a helpful assistant with tool calling capabilities.
+    {{- end }}<|eot_id|>
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 }}
+    {{- if eq .Role "user" }}<|start_header_id|>user<|end_header_id|>
+    {{- if and $.Tools $last }}
+    Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt.
+    Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. Do not use variables.
+    {{ range $.Tools }}
+    {{- . }}
+    {{ end }}
+    {{ .Content }}<|eot_id|>
+    {{- else }}
+    {{ .Content }}<|eot_id|>
+    {{- end }}{{ if $last }}<|start_header_id|>assistant<|end_header_id|>
+    {{ end }}
+    {{- else if eq .Role "assistant" }}<|start_header_id|>assistant<|end_header_id|>
+    {{- if .ToolCalls }}
+    {{ range .ToolCalls }}
+    {"name": "{{ .Function.Name }}", "parameters": {{ .Function.Arguments }}}{{ end }}
+    {{- else }}
+    {{ .Content }}
+    {{- end }}{{ if not $last }}<|eot_id|>{{ end }}
+    {{- else if eq .Role "tool" }}<|start_header_id|>ipython<|end_header_id|>
+    {{ .Content }}<|eot_id|>{{ if $last }}<|start_header_id|>assistant<|end_header_id|>
+    {{ end }}
+    {{- end }}
+    {{- end }}""",
+                "stop_tokens": ["<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"]
+            },
+            "qwen": {
+                "template": """{{- if .Suffix }}<|fim_prefix|>{{ .Prompt }}<|fim_suffix|>{{ .Suffix }}<|fim_middle|>
+    {{- else if .Messages }}
+    {{- if or .System .Tools }}<|im_start|>system
+    {{- if .System }}
+    {{ .System }}
+    {{- end }}
+    {{- if .Tools }}
+    # Tools
+    You may call one or more functions to assist with the user query.
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {{- range .Tools }}
+    {"type": "function", "function": {{ .Function }}}
+    {{- end }}
+    </tools>
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    {{- end }}<|im_end|>
+    {{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    {{- if eq .Role "user" }}<|im_start|>user
+    {{ .Content }}<|im_end|>
+    {{ else if eq .Role "assistant" }}<|im_start|>assistant
+    {{ if .Content }}{{ .Content }}
+    {{- else if .ToolCalls }}<tool_call>
+    {{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
+    {{ end }}</tool_call>
+    {{- end }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- else if eq .Role "tool" }}<|im_start|>user
+    <tool_response>
+    {{ .Content }}
+    </tool_response><|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
+    {{ end }}
+    {{- end }}
+    {{- else }}
+    {{- if .System }}<|im_start|>system
+    {{ .System }}<|im_end|>
+    {{ end }}{{ if .Prompt }}<|im_start|>user
+    {{ .Prompt }}<|im_end|>
+    {{ end }}<|im_start|>assistant
+    {{ end }}{{ .Response }}{{ if .Response }}<|im_end|>{{ end }}""",
+                "system": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
+                "num_ctx": 32768,
+                "stop_tokens": ["<|endoftext|>"]
+            },
+            "mistral": {
+                "template": "[INST] {{ if .System }}{{ .System }} {{ end }}{{ .Prompt }} [/INST]",
+                "stop_tokens": ["[INST]", "[/INST]"]
+            },
+            "phi": {
+                "template": """{{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    <|im_start|>{{ .Role }}<|im_sep|>
+    {{ .Content }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_end|>
+    <|im_start|>assistant<|im_sep|>
+    {{ end }}
+    {{- end }}""",
+                "stop_tokens": ["<|im_start|>", "<|im_end|>", "<|im_sep|>"]
+            },
+            "deepseek": {
+                "template": """{{- if .System }}{{ .System }}{{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1}}
+    {{- if eq .Role "user" }}
+    {{ .Content }}
+    {{- else if eq .Role "assistant" }}
+    {{ .Content }}{{- if not $last }}
+    {{- end }}
+    {{- end }}
+    {{- if and $last (ne .Role "assistant") }}
+    {{ end }}
+    {{- end }}""",
+                "stop_tokens": ["", "", "", ""]
+            },
+            "llava": {
+                "template": """{{- if .Suffix }}<|fim_prefix|>{{ .Prompt }}<|fim_suffix|>{{ .Suffix }}<|fim_middle|>
+    {{- else if .Messages }}
+    {{- if or .System .Tools }}<|im_start|>system
+    {{- if .System }}
+    {{ .System }}
+    {{- end }}
+    {{- if .Tools }}
+    # Tools
+    You may call one or more functions to assist with the user query.
+    You are provided with function signatures within <tools></tools> XML tags:
+    <tools>
+    {{- range .Tools }}
+    {"type": "function", "function": {{ .Function }}}
+    {{- end }}
+    </tools>
+    For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+    <tool_call>
+    {"name": <function-name>, "arguments": <args-json-object>}
+    </tool_call>
+    {{- end }}<|im_end|>
+    {{ end }}
+    {{- range $i, $_ := .Messages }}
+    {{- $last := eq (len (slice $.Messages $i)) 1 -}}
+    {{- if eq .Role "user" }}<|im_start|>user
+    {{ .Content }}<|im_end|>
+    {{ else if eq .Role "assistant" }}<|im_start|>assistant
+    {{ if .Content }}{{ .Content }}
+    {{- else if .ToolCalls }}<tool_call>
+    {{ range .ToolCalls }}{"name": "{{ .Function.Name }}", "arguments": {{ .Function.Arguments }}}
+    {{ end }}</tool_call>
+    {{- end }}{{ if not $last }}<|im_end|>
+    {{ end }}
+    {{- else if eq .Role "tool" }}<|im_start|>user
+    <tool_response>
+    {{ .Content }}
+    </tool_response><|im_end|>
+    {{ end }}
+    {{- if and (ne .Role "assistant") $last }}<|im_start|>assistant
+    {{ end }}
+    {{- end }}
+    {{- else }}
+    {{- if .System }}<|im_start|>system
+    {{ .System }}<|im_end|>
+    {{ end }}{{ if .Prompt }}<|im_start|>user
+    {{ .Prompt }}<|im_end|>
+    {{ end }}<|im_start|>assistant
+    {{ end }}{{ .Response }}{{ if .Response }}<|im_end|>{{ end }}""",
+                "stop_tokens": ["</s>", "USER:", "ASSSISTANT:"]
+            }
+        }
+        # Select mapping by checking if any key is in the model_name.
+        chosen = None
+        for key, settings in mapping.items():
+            if key in model_name:
+                chosen = settings
+                break
+        if chosen is None:
+            # Fallback default
+            chosen = {
+                "template": """{{ if .System }}<|start_header_id|>system<|end_header_id|>
+    {{ .System }}<|eot_id|>{{ end }}{{ if .Prompt }}<|start_header_id|>user<|end_header_id|>
+    {{ .Prompt }}<|eot_id|>{{ end }}<|start_header_id|>assistant<|end_header_id|>
+    {{ .Response }}<|eot_id|>""",
+                "stop_tokens": ["<|start_header_id|>", "<|end_header_id|>", "<|eot_id|>"]
+            }
+        # Build the stop parameter lines.
+        stop_params = "\n".join([f"PARAMETER stop {token}" for token in chosen["stop_tokens"]])
+        # Optionally include a SYSTEM line and num_ctx if defined in the mapping.
+        system_line = ""
+        if "system" in chosen:
+            system_line = f"SYSTEM {chosen['system']}\n"
+        num_ctx_line = ""
+        if "num_ctx" in chosen:
+            num_ctx_line = f"PARAMETER num_ctx {chosen['num_ctx']}\n"
+        # Assemble and return the modelfile content.
+        return f"""FROM {output_model}
+    TEMPLATE \"\"\"{chosen['template']}\"\"\"
+    {system_line}{num_ctx_line}{stop_params}
+    """

    def create_and_push_ollama_model(self):
        modelfile_content = self.prepare_modelfile_content()
@@ -0,0 +1,234 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+This script finetunes a vision language model using Unsloth's fast training framework.
+It supports vision tasks by converting raw image-caption samples into a conversation format, 
+adding vision-specific LoRA adapters, and training using TRL's SFTTrainer with UnslothVisionDataCollator.
+"""
+
+import os
+import sys
+import yaml
+import torch
+import shutil
+import subprocess
+
+from datasets import load_dataset, concatenate_datasets
+from unsloth import FastVisionModel, is_bf16_supported
+from unsloth.trainer import UnslothVisionDataCollator
+from trl import SFTTrainer, SFTConfig
+
+
+class TrainVisionModel:
+    def __init__(self, config_path="config.yaml"):
+        self.load_config(config_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = None
+        self.hf_tokenizer = None  # The underlying tokenizer
+
+    def load_config(self, path):
+        with open(path, "r") as file:
+            self.config = yaml.safe_load(file)
+        print("DEBUG: Loaded config:", self.config)
+
+    def print_system_info(self):
+        print("DEBUG: PyTorch version:", torch.__version__)
+        print("DEBUG: CUDA version:", torch.version.cuda)
+        if torch.cuda.is_available():
+            print("DEBUG: CUDA Device Capability:", torch.cuda.get_device_capability())
+        else:
+            print("DEBUG: CUDA is not available")
+        print("DEBUG: Python Version:", sys.version)
+        print("DEBUG: Python Path:", sys.executable)
+
+    def check_gpu(self):
+        gpu_stats = torch.cuda.get_device_properties(0)
+        print(f"DEBUG: GPU = {gpu_stats.name}. Max memory = {round(gpu_stats.total_memory/(1024**3),3)} GB.")
+
+    def check_ram(self):
+        from psutil import virtual_memory
+        ram_gb = virtual_memory().total / 1e9
+        print(f"DEBUG: Your runtime has {ram_gb:.1f} gigabytes of available RAM")
+        if ram_gb < 20:
+            print("DEBUG: Not using a high-RAM runtime")
+        else:
+            print("DEBUG: You are using a high-RAM runtime!")
+
+    def prepare_model(self):
+        print("DEBUG: Preparing vision model and tokenizer...")
+        self.model, original_tokenizer = FastVisionModel.from_pretrained(
+            model_name=self.config["model_name"],
+            load_in_4bit=self.config["load_in_4bit"],
+            use_gradient_checkpointing="unsloth"
+        )
+        print("DEBUG: Vision model and original tokenizer loaded.")
+        if original_tokenizer.pad_token is None:
+            original_tokenizer.pad_token = original_tokenizer.eos_token
+        original_tokenizer.model_max_length = self.config.get("max_seq_length", 2048)
+        self.hf_tokenizer = original_tokenizer
+        
+        # Add vision-specific LoRA adapters
+        self.model = FastVisionModel.get_peft_model(
+            self.model,
+            finetune_vision_layers=self.config.get("finetune_vision_layers", False),
+            finetune_language_layers=self.config.get("finetune_language_layers", True),
+            finetune_attention_modules=self.config.get("finetune_attention_modules", True),
+            finetune_mlp_modules=self.config.get("finetune_mlp_modules", True),
+            r=16,
+            lora_alpha=16,
+            lora_dropout=0,
+            bias="none",
+            random_state=3407,
+            use_rslora=False,
+            loftq_config=None
+        )
+        print("DEBUG: Vision LoRA adapters added.")
+
+    def convert_sample(self, sample):
+        # Use a default instruction or one from config
+        instr = self.config.get("vision_instruction", "You are an expert radiographer. Describe accurately what you see in this image.")
+        conversation = [
+            {"role": "user", "content": [
+                {"type": "text", "text": instr},
+                {"type": "image", "image": sample["image"]}
+            ]},
+            {"role": "assistant", "content": [
+                {"type": "text", "text": sample["caption"]}
+            ]}
+        ]
+        return {"messages": conversation}
+
+    def load_datasets(self):
+        datasets = []
+        for dataset_info in self.config["dataset"]:
+            print("DEBUG: Loading vision dataset:", dataset_info)
+            ds = load_dataset(dataset_info["name"], split=dataset_info.get("split_type", "train"))
+            print("DEBUG: Converting dataset to vision conversation format...")
+            ds = ds.map(self.convert_sample)
+            datasets.append(ds)
+        combined = concatenate_datasets(datasets)
+        print("DEBUG: Combined vision dataset has", len(combined), "examples.")
+        return combined
+
+    def train_model(self):
+        print("DEBUG: Starting vision training...")
+        raw_dataset = self.load_datasets()
+        
+        # Build training arguments using SFTConfig for vision tasks
+        sft_config = SFTConfig(
+            per_device_train_batch_size=self.config.get("per_device_train_batch_size", 2),
+            gradient_accumulation_steps=self.config.get("gradient_accumulation_steps", 4),
+            warmup_steps=self.config.get("warmup_steps", 5),
+            max_steps=self.config.get("max_steps", 30),
+            learning_rate=self.config.get("learning_rate", 2e-4),
+            fp16=self.config.get("fp16", not is_bf16_supported()),
+            bf16=self.config.get("bf16", is_bf16_supported()),
+            logging_steps=self.config.get("logging_steps", 1),
+            optim=self.config.get("optim", "adamw_8bit"),
+            weight_decay=self.config.get("weight_decay", 0.01),
+            lr_scheduler_type=self.config.get("lr_scheduler_type", "linear"),
+            seed=self.config.get("seed", 3407),
+            output_dir=self.config.get("output_dir", "outputs"),
+            report_to="none" if not os.getenv("PRAISON_WANDB") else "wandb",
+            remove_unused_columns=False,
+            dataset_text_field="",
+            dataset_kwargs={"skip_prepare_dataset": True},
+            dataset_num_proc=self.config.get("dataset_num_proc", 4),
+            max_seq_length=self.config.get("max_seq_length", 2048)
+        )
+        
+        trainer = SFTTrainer(
+            model=self.model,
+            tokenizer=self.hf_tokenizer,
+            data_collator=UnslothVisionDataCollator(self.model, self.hf_tokenizer),
+            train_dataset=raw_dataset,
+            args=sft_config
+        )
+        print("DEBUG: Beginning vision trainer.train() ...")
+        trainer.train()
+        print("DEBUG: Vision training complete. Saving model and tokenizer locally...")
+        self.model.save_pretrained("lora_vision_model")
+        self.hf_tokenizer.save_pretrained("lora_vision_model")
+        print("DEBUG: Saved vision model and tokenizer to 'lora_vision_model'.")
+
+    def vision_inference(self, instruction, image):
+        FastVisionModel.for_inference(self.model)
+        messages = [
+            {"role": "user", "content": [
+                {"type": "image"},
+                {"type": "text", "text": instruction}
+            ]}
+        ]
+        input_text = self.hf_tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+        inputs = self.hf_tokenizer(
+            image,
+            input_text,
+            add_special_tokens=False,
+            return_tensors="pt"
+        ).to("cuda")
+        outputs = self.model.generate(
+            **inputs,
+            max_new_tokens=128,
+            use_cache=True,
+            temperature=1.5,
+            min_p=0.1
+        )
+        print("DEBUG: Vision inference output:", self.hf_tokenizer.batch_decode(outputs))
+
+    def save_model_merged(self):
+        if os.path.exists(self.config["hf_model_name"]):
+            shutil.rmtree(self.config["hf_model_name"])
+        self.model.push_to_hub_merged(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            save_method="merged_16bit",
+            token=os.getenv("HF_TOKEN")
+        )
+
+    def push_model_gguf(self):
+        self.model.push_to_hub_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method=self.config.get("quantization_method", "q4_k_m"),
+            token=os.getenv("HF_TOKEN")
+        )
+
+    def save_model_gguf(self):
+        self.model.save_pretrained_gguf(
+            self.config["hf_model_name"],
+            self.hf_tokenizer,
+            quantization_method="q4_k_m"
+        )
+
+    def run(self):
+        self.print_system_info()
+        self.check_gpu()
+        self.check_ram()
+        if self.config.get("train", "true").lower() == "true":
+            self.prepare_model()
+            self.train_model()
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="PraisonAI Vision Training Script")
+    parser.add_argument("command", choices=["train", "inference"], help="Command to execute")
+    parser.add_argument("--config", default="config.yaml", help="Path to configuration file")
+    args = parser.parse_args()
+
+    trainer_obj = TrainVisionModel(config_path=args.config)
+    if args.command == "train":
+        trainer_obj.run()
+    elif args.command == "inference":
+        # For inference, we load a sample image from the first dataset
+        instr = trainer_obj.config.get("vision_instruction", "You are an expert radiographer. Describe accurately what you see in this image.")
+        ds_info = trainer_obj.config["dataset"][0]
+        ds = load_dataset(ds_info["name"], split=ds_info.get("split_type", "train"))
+        sample_image = ds[0]["image"]
+        if trainer_obj.model is None or trainer_obj.hf_tokenizer is None:
+            trainer_obj.prepare_model()
+        trainer_obj.vision_inference(instr, sample_image)
+
+
+if __name__ == "__main__":
+    main() 
@@ -1,6 +1,6 @@
 [project]
 name = "PraisonAI"
-version = "2.0.66"
+version = "2.0.75"
 description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human-agent collaboration."
 readme = "README.md"
 license = ""
@@ -84,7 +84,7 @@ autogen = ["pyautogen>=0.2.19", "praisonai-tools>=0.0.7", "crewai"]

 [tool.poetry]
 name = "PraisonAI"
-version = "2.0.66"
+version = "2.0.75"
 description = "PraisonAI is an AI Agents Framework with Self Reflection. PraisonAI application combines PraisonAI Agents, AutoGen, and CrewAI into a low-code solution for building and managing multi-agent LLM systems, focusing on simplicity, customisation, and efficient human–agent collaboration."
 authors = ["Mervin Praison"]
 license = ""
@@ -3060,7 +3060,7 @@ wheels = [

 [[package]]
 name = "praisonai"
-version = "2.0.66"
+version = "2.0.75"
 source = { editable = "." }
 dependencies = [
    { name = "instructor" },
Author	SHA1	Message	Date
Mervin Praison	11db6e18ba	Merge pull request #382 from MervinPraison/develop Add Vision Model Training Support and Update Unsloth Dependency	2025-02-09 00:31:40 +00:00
MervinPraison	0d6f698e6c	Add Vision Model Training Support and Update Unsloth Dependency - Enhanced CLI to dynamically select training script based on model type (vision vs. standard) - Added support for detecting vision models through flexible name matching - Updated Unsloth dependency to latest commit - Bumped version to 2.0.75 across all relevant files	2025-02-09 00:31:20 +00:00
Mervin Praison	b8574f78e2	Merge pull request #381 from MervinPraison/develop Bump version to 2.0.74 and refactor train.py template generation	2025-02-07 13:14:32 +00:00
MervinPraison	e631774a0f	Bump version to 2.0.74 and refactor train.py template generation - Updated version to 2.0.74 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Simplified training template generation in train.py - Removed complex stop tokens and simplified template structure - Dynamically built training arguments with improved configuration flexibility	2025-02-07 13:14:16 +00:00
Mervin Praison	df735aa748	Merge pull request #378 from MervinPraison/develop Enhance train.py with Configurable Training Parameters	2025-02-06 07:48:00 +00:00
MervinPraison	3bdf5b6c65	Enhance train.py with Configurable Training Parameters - Updated version to 2.0.73 across multiple files - Refactored TrainModel to use configurable training parameters with default fallbacks - Improved flexibility of training configuration by allowing custom settings via config - Dynamically set training arguments based on provided configuration	2025-02-06 07:47:40 +00:00
Mervin Praison	361fa67108	Merge pull request #377 from MervinPraison/develop Bump version to 2.0.72 and enhance CLI configuration in train and cli…	2025-02-06 07:26:59 +00:00
MervinPraison	c5abcf362a	Bump version to 2.0.72 and enhance CLI configuration in train and cli modules - Updated version to 2.0.72 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Refined CLI configuration handling in cli.py for config.yaml generation - Simplified train.py main function to always use provided config path - Improved configuration flexibility for training and model setup	2025-02-06 07:26:41 +00:00
Mervin Praison	7ed2d99a11	Merge pull request #376 from MervinPraison/develop Enhance train.py CLI configuration handling and bump version to 2.0.71	2025-02-06 07:19:11 +00:00
MervinPraison	81307da7c8	Enhance train.py CLI configuration handling and bump version to 2.0.71 - Updated version to 2.0.71 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Modified train.py to use existing config.yaml when no CLI arguments are provided - Improved CLI configuration flexibility for training command	2025-02-06 07:18:56 +00:00
Mervin Praison	fb75153b5c	Merge pull request #373 from MervinPraison/develop Bump version to 2.0.70 and update configuration handling	2025-02-05 06:28:07 +00:00
MervinPraison	cddcfcf170	Bump version to 2.0.70 and update configuration handling - Updated version to 2.0.70 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Enhanced CLI configuration handling in cli.py to load existing config when not creating a new one - Minor improvements to configuration management and deployment scripts	2025-02-05 06:27:50 +00:00
Mervin Praison	398b0473e0	Merge pull request #372 from MervinPraison/develop Update train.mdx with Wandb configuration and requirements	2025-02-05 05:07:23 +00:00
MervinPraison	c5d260782a	Update train.mdx with Wandb configuration and requirements - Added detailed requirements for training, including Huggingface token and model details - Included Wandb login and configuration instructions - Updated documentation with environment variables for Wandb logging	2025-02-05 05:07:08 +00:00
Mervin Praison	116ac5a64d	Merge pull request #370 from MervinPraison/develop Add Wandb Support for Training and Bump Version to 2.0.69	2025-02-04 11:18:42 +00:00
MervinPraison	abf6e2e8d1	Add Wandb Support for Training and Bump Version to 2.0.69 - Updated version to 2.0.69 across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Enhanced train.py with optional Wandb logging support - Added environment variable-based Wandb configuration - Implemented dynamic Wandb reporting and run naming	2025-02-04 11:18:20 +00:00
Mervin Praison	cb2b3617fa	Merge pull request #369 from MervinPraison/develop Bump version to 2.0.68 and enhance Ollama model template generation	2025-02-04 11:10:50 +00:00
MervinPraison	6de67c7d0d	Bump version to 2.0.68 and enhance Ollama model template generation - Updated version across multiple files (praisonai.rb, pyproject.toml, uv.lock, Dockerfile, docs) - Significantly improved train.py's prepare_modelfile_content() method - Added comprehensive template and stop token mapping for multiple model architectures - Implemented dynamic template selection based on model name - Enhanced Ollama model file generation with more flexible configuration	2025-02-04 11:10:26 +00:00
Mervin Praison	5699d90894	Merge pull request #368 from MervinPraison/develop Enhance train.py with Dynamic Stop Token Generation for Ollama Models	2025-02-04 10:33:19 +00:00
MervinPraison	2982e8ea9c	Enhance train.py with Dynamic Stop Token Generation for Ollama Models - Added intelligent stop token inference for various model architectures in train.py - Implemented dynamic stop token selection based on model name - Supported models include Phi, LLaVA, Mistral, Qwen, DeepSeek, and default models - Updated Modelfile generation to include model-specific stop tokens - Improved Ollama model preparation flexibility	2025-02-04 10:33:05 +00:00