-
- -- **load image** - - ```python - import cv2 - - image = cv2.imread("...") - ``` - -- **create and refine marks** - - ```python - import maestro - - generator = maestro.SegmentAnythingMarkGenerator(device='cuda') - marks = generator.generate(image=image) - marks = maestro.refine_marks(marks=marks) - ``` - -- **visualize marks** - - ```python - mark_visualizer = maestro.MarkVisualizer() - marked_image = mark_visualizer.visualize(image=image, marks=marks) - ``` -  - -- **prompt** - - ```python - prompt = "Find dog." - - response = maestro.prompt_image(api_key=api_key, image=marked_image, prompt=prompt) - ``` - - ``` - >>> "The dog is prominently featured in the center of the image with the label [9]." - ``` - -- **extract related marks** - - ```python - masks = maestro.extract_relevant_masks(text=response, detections=refined_marks) - ``` - - ``` - >>> {'6': array([ - ... [False, False, False, ..., False, False, False], - ... [False, False, False, ..., False, False, False], - ... [False, False, False, ..., False, False, False], - ... ..., - ... [ True, True, True, ..., False, False, False], - ... [ True, True, True, ..., False, False, False], - ... [ True, True, True, ..., False, False, False]]) - ... } - ``` - -
-
-
-
-## 🚧 roadmap
-
-- [ ] Rewriting the `maestro` API.
-- [ ] Update [HF space](https://huggingface.co/spaces/Roboflow/SoM).
-- [ ] Documentation page.
-- [ ] Add GroundingDINO prompting strategy.
-- [ ] CovVLM demo.
-- [ ] Qwen-VL demo.
-
-## 💜 acknowledgement
-
-- [Set-of-Mark Prompting Unleashes Extraordinary Visual Grounding
-in GPT-4V](https://arxiv.org/abs/2310.11441) by Jianwei Yang, Hao Zhang, Feng Li, Xueyan
-Zou, Chunyuan Li, Jianfeng Gao.
-- [The Dawn of LMMs: Preliminary Explorations with GPT-4V(ision)](https://arxiv.org/abs/2309.17421)
-by Zhengyuan Yang, Linjie Li, Kevin Lin, Jianfeng Wang, Chung-Ching Lin, Zicheng Liu,
-Lijuan Wang
-
## 🦸 contribution
-We would love your help in making this repository even better! If you noticed any bug,
-or if you have any suggestions for improvement, feel free to open an
+We would love your help in making this repository even better! We are especially
+looking for contributors with experience in fine-tuning vision-language models (VLMs).
+If you notice any bugs or have suggestions for improvement, feel free to open an
[issue](https://github.com/roboflow/multimodal-maestro/issues) or submit a
[pull request](https://github.com/roboflow/multimodal-maestro/pulls).
diff --git a/maestro/cli/__init__.py b/maestro/cli/__init__.py
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/maestro/cli/__init__.py
@@ -0,0 +1 @@
+
diff --git a/maestro/cli/env.py b/maestro/cli/env.py
new file mode 100644
index 0000000..b95525e
--- /dev/null
+++ b/maestro/cli/env.py
@@ -0,0 +1,2 @@
+DISABLE_RECIPE_IMPORTS_WARNINGS_ENV = "DISABLE_RECIPE_IMPORTS_WARNINGS"
+DEFAULT_DISABLE_RECIPE_IMPORTS_WARNINGS_ENV = "False"
diff --git a/maestro/cli/introspection.py b/maestro/cli/introspection.py
new file mode 100644
index 0000000..fc6aef9
--- /dev/null
+++ b/maestro/cli/introspection.py
@@ -0,0 +1,37 @@
+import os
+
+import typer
+
+from maestro.cli.env import DISABLE_RECIPE_IMPORTS_WARNINGS_ENV, \
+ DEFAULT_DISABLE_RECIPE_IMPORTS_WARNINGS_ENV
+from maestro.cli.utils import str2bool
+
+
+def find_training_recipes(app: typer.Typer) -> None:
+ try:
+ from maestro.trainer.models.florence_2.entrypoint import florence_2_app
+
+ app.add_typer(florence_2_app, name="florence2")
+ except Exception:
+ _warn_about_recipe_import_error(model_name="Florence 2")
+
+ try:
+ from maestro.trainer.models.paligemma.entrypoint import paligemma_app
+
+ app.add_typer(paligemma_app, name="paligemma")
+ except Exception:
+ _warn_about_recipe_import_error(model_name="PaliGemma")
+
+
+def _warn_about_recipe_import_error(model_name: str) -> None:
+ disable_warnings = str2bool(
+ os.getenv(
+ DISABLE_RECIPE_IMPORTS_WARNINGS_ENV,
+ DEFAULT_DISABLE_RECIPE_IMPORTS_WARNINGS_ENV,
+ )
+ )
+ if disable_warnings:
+ return None
+ warning = typer.style("WARNING", fg=typer.colors.RED, bold=True)
+ message = "🚧 " + warning + f" cannot import recipe for {model_name}"
+ typer.echo(message)
diff --git a/maestro/cli/main.py b/maestro/cli/main.py
new file mode 100644
index 0000000..b600e3a
--- /dev/null
+++ b/maestro/cli/main.py
@@ -0,0 +1,15 @@
+import typer
+
+from maestro.cli.introspection import find_training_recipes
+
+app = typer.Typer()
+find_training_recipes(app=app)
+
+
+@app.command(help="Display information about maestro")
+def info():
+ typer.echo("Welcome to maestro CLI. Let's train some VLM! 🏋")
+
+
+if __name__ == "__main__":
+ app()
diff --git a/maestro/cli/utils.py b/maestro/cli/utils.py
new file mode 100644
index 0000000..0751fef
--- /dev/null
+++ b/maestro/cli/utils.py
@@ -0,0 +1,2 @@
+def str2bool(value: str) -> bool:
+ return value.lower() in {"y", "t", "yes", "true"}
diff --git a/maestro/trainer/__init__.py b/maestro/trainer/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/maestro/trainer/common/__init__.py b/maestro/trainer/common/__init__.py
new file mode 100644
index 0000000..0bfe6b2
--- /dev/null
+++ b/maestro/trainer/common/__init__.py
@@ -0,0 +1 @@
+from maestro.trainer.common.utils.metrics import MeanAveragePrecisionMetric
diff --git a/maestro/trainer/common/configuration/__init__.py b/maestro/trainer/common/configuration/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/maestro/trainer/common/configuration/env.py b/maestro/trainer/common/configuration/env.py
new file mode 100644
index 0000000..eeb9105
--- /dev/null
+++ b/maestro/trainer/common/configuration/env.py
@@ -0,0 +1,5 @@
+SEED_ENV = "SEED"
+DEFAULT_SEED = "42"
+CUDA_DEVICE_ENV = "CUDA_DEVICE"
+DEFAULT_CUDA_DEVICE = "cuda:0"
+HF_TOKEN_ENV = "HF_TOKEN"
diff --git a/maestro/trainer/common/data_loaders/__init__.py b/maestro/trainer/common/data_loaders/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/maestro/trainer/common/data_loaders/datasets.py b/maestro/trainer/common/data_loaders/datasets.py
new file mode 100644
index 0000000..b79cfec
--- /dev/null
+++ b/maestro/trainer/common/data_loaders/datasets.py
@@ -0,0 +1,50 @@
+import json
+import os
+from typing import List, Dict, Any, Tuple
+
+from PIL import Image
+from transformers.pipelines.base import Dataset
+
+
+class JSONLDataset:
+ def __init__(self, jsonl_file_path: str, image_directory_path: str):
+ self.jsonl_file_path = jsonl_file_path
+ self.image_directory_path = image_directory_path
+ self.entries = self._load_entries()
+
+ def _load_entries(self) -> List[Dict[str, Any]]:
+ entries = []
+ with open(self.jsonl_file_path, "r") as file:
+ for line in file:
+ data = json.loads(line)
+ entries.append(data)
+ return entries
+
+ def __len__(self) -> int:
+ return len(self.entries)
+
+ def __getitem__(self, idx: int) -> Tuple[Image.Image, Dict[str, Any]]:
+ if idx < 0 or idx >= len(self.entries):
+ raise IndexError("Index out of range")
+
+ entry = self.entries[idx]
+ image_path = os.path.join(self.image_directory_path, entry["image"])
+ try:
+ image = Image.open(image_path)
+ return (image, entry)
+ except FileNotFoundError:
+ raise FileNotFoundError(f"Image file {image_path} not found.")
+
+
+class DetectionDataset(Dataset):
+ def __init__(self, jsonl_file_path: str, image_directory_path: str):
+ self.dataset = JSONLDataset(jsonl_file_path, image_directory_path)
+
+ def __len__(self):
+ return len(self.dataset)
+
+ def __getitem__(self, idx):
+ image, data = self.dataset[idx]
+ prefix = data["prefix"]
+ suffix = data["suffix"]
+ return prefix, suffix, image
diff --git a/maestro/trainer/common/data_loaders/jsonl.py b/maestro/trainer/common/data_loaders/jsonl.py
new file mode 100644
index 0000000..3630e11
--- /dev/null
+++ b/maestro/trainer/common/data_loaders/jsonl.py
@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import random
+from typing import List
+
+from torch.utils.data import Dataset
+
+from maestro.trainer.common.utils.file_system import read_jsonl
+
+
+class JSONLDataset(Dataset):
+ # TODO: implementation could be better - avoiding loading
+ # whole files to memory
+
+ @classmethod
+ def from_jsonl_file(cls, path: str) -> JSONLDataset:
+ file_content = read_jsonl(path=path)
+ random.shuffle(file_content)
+ return cls(jsons=file_content)
+
+ def __init__(self, jsons: List[dict]):
+ self.jsons = jsons
+
+ def __getitem__(self, index):
+ return self.jsons[index]
+
+ def __len__(self):
+ return len(self.jsons)
+
+ def shuffle(self):
+ random.shuffle(self.jsons)
diff --git a/maestro/trainer/common/utils/__init__.py b/maestro/trainer/common/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/maestro/trainer/common/utils/file_system.py b/maestro/trainer/common/utils/file_system.py
new file mode 100644
index 0000000..da8e5a9
--- /dev/null
+++ b/maestro/trainer/common/utils/file_system.py
@@ -0,0 +1,59 @@
+import json
+import os
+from glob import glob
+from typing import Union, List
+
+
+def read_jsonl(path: str) -> List[dict]:
+ file_lines = read_file(
+ path=path,
+ split_lines=True,
+ )
+ return [json.loads(line) for line in file_lines]
+
+
+def read_file(
+ path: str,
+ split_lines: bool = False,
+ strip_white_spaces: bool = False,
+ line_separator: str = "\n",
+) -> Union[str, List[str]]:
+ with open(path, "r") as f:
+ file_content = f.read()
+ if strip_white_spaces:
+ file_content = file_content.strip()
+ if not split_lines:
+ return file_content
+ lines = file_content.split(line_separator)
+ if not strip_white_spaces:
+ return lines
+ return [line.strip() for line in lines]
+
+
+def save_json(path: str, content: dict) -> None:
+ ensure_parent_dir_exists(path=path)
+ with open(path, "w") as f:
+ json.dump(content, f, indent=4)
+
+
+def ensure_parent_dir_exists(path: str) -> None:
+ parent_dir = os.path.dirname(os.path.abspath(path))
+ os.makedirs(parent_dir, exist_ok=True)
+
+
+def create_new_run_directory(base_output_dir: str) -> str:
+ """
+ Creates a new numbered directory for the current training run.
+
+ Args:
+ base_output_dir (str): The base directory where all run directories are stored.
+
+ Returns:
+ str: The path to the newly created run directory.
+ """
+ base_output_dir = os.path.abspath(base_output_dir)
+ existing_run_dirs = [d for d in glob(os.path.join(base_output_dir, "*")) if os.path.isdir(d)]
+ new_run_number = len(existing_run_dirs) + 1
+ new_run_dir = os.path.join(base_output_dir, str(new_run_number))
+ os.makedirs(new_run_dir, exist_ok=True)
+ return new_run_dir
diff --git a/maestro/trainer/common/utils/leaderboard.py b/maestro/trainer/common/utils/leaderboard.py
new file mode 100644
index 0000000..6f381fe
--- /dev/null
+++ b/maestro/trainer/common/utils/leaderboard.py
@@ -0,0 +1,42 @@
+from typing import Dict, Tuple, Optional
+
+
+class CheckpointsLeaderboard:
+
+ def __init__(
+ self,
+ max_checkpoints: int,
+ ):
+ self._max_checkpoints = max(max_checkpoints, 1)
+ self._leaderboard: Dict[int, Tuple[str, float]] = {}
+
+ def register_checkpoint(self, epoch: int, path: str, loss: float) -> Tuple[bool, Optional[str]]:
+ if len(self._leaderboard) < self._max_checkpoints:
+ self._leaderboard[epoch] = (path, loss)
+ return True, None
+ max_loss_key, max_loss_in_leaderboard = None, None
+ for key, (_, loss) in self._leaderboard.items():
+ if max_loss_in_leaderboard is None:
+ max_loss_key = key
+ max_loss_in_leaderboard = loss
+ if loss > max_loss_in_leaderboard: # type: ignore
+ max_loss_key = key
+ max_loss_in_leaderboard = loss
+ if loss >= max_loss_in_leaderboard: # type: ignore
+ return False, None
+ to_be_removed, _ = self._leaderboard.pop(max_loss_key) # type: ignore
+ self._leaderboard[epoch] = (path, loss)
+ return True, to_be_removed
+
+ def get_best_model(self) -> str:
+ min_loss_key, min_loss_in_leaderboard = None, None
+ for key, (_, loss) in self._leaderboard.items():
+ if min_loss_in_leaderboard is None:
+ min_loss_key = key
+ min_loss_in_leaderboard = loss
+ if loss < min_loss_in_leaderboard: # type: ignore
+ min_loss_key = key
+ min_loss_in_leaderboard = loss
+ if min_loss_key is None:
+ raise RuntimeError("Could not retrieve best model")
+ return self._leaderboard[min_loss_key][0]
diff --git a/maestro/trainer/common/utils/metrics.py b/maestro/trainer/common/utils/metrics.py
new file mode 100644
index 0000000..0d1fb9b
--- /dev/null
+++ b/maestro/trainer/common/utils/metrics.py
@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import base64
+import html
+import io
+import json
+import os
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Any, Dict, List, Tuple
+
+import matplotlib.pyplot as plt
+import supervision as sv
+from PIL import Image
+from supervision.metrics.mean_average_precision import MeanAveragePrecision
+
+
+class BaseMetric(ABC):
+ """
+ Abstract base class for custom metrics. Subclasses must implement
+ the 'describe' and 'compute' methods.
+ """
+
+ @abstractmethod
+ def describe(self) -> List[str]:
+ """
+ Describe the names of the metrics that this class will compute.
+
+ Returns:
+ List[str]: A list of metric names that will be computed.
+ """
+ pass
+
+ @abstractmethod
+ def compute(self, targets: List[Any], predictions: List[Any]) -> Dict[str, float]:
+ """
+ Compute the metric based on the targets and predictions.
+
+ Args:
+ targets (List[Any]): The ground truth.
+ predictions (List[Any]): The prediction result.
+
+ Returns:
+ Dict[str, float]: A dictionary of computed metrics with metric names as
+ keys and their values.
+ """
+ pass
+
+
+class MeanAveragePrecisionMetric(BaseMetric):
+ """
+ A class used to compute the Mean Average Precision (mAP) metric.
+ """
+
+ def describe(self) -> List[str]:
+ """
+ Returns a list of metric names that this class will compute.
+
+ Returns:
+ List[str]: A list of metric names.
+ """
+ return ["map50:95", "map50", "map75"]
+
+ def compute(
+ self,
+ targets: List[sv.Detections],
+ predictions: List[sv.Detections]
+ ) -> Dict[str, float]:
+ """
+ Computes the mAP metrics based on the targets and predictions.
+
+ Args:
+ targets (List[sv.Detections]): The ground truth detections.
+ predictions (List[sv.Detections]): The predicted detections.
+
+ Returns:
+ Dict[str, float]: A dictionary of computed mAP metrics with metric names as
+ keys and their values.
+ """
+ result = MeanAveragePrecision().update(
+ targets=targets, predictions=predictions).compute()
+ return {
+ "map50:95": result.map50_95,
+ "map50": result.map50,
+ "map75": result.map75
+ }
+
+
+class MetricsTracker:
+
+ @classmethod
+ def init(cls, metrics: List[str]) -> MetricsTracker:
+ return cls(metrics={metric: [] for metric in metrics})
+
+ def __init__(self, metrics: Dict[str, List[Tuple[int, int, float]]]):
+ self._metrics = metrics
+
+ def register(self, metric: str, epoch: int, step: int, value: float) -> None:
+ self._metrics[metric].append((epoch, step, value))
+
+ def describe_metrics(self) -> List[str]:
+ return list(self._metrics.keys())
+
+ def get_metric_values(
+ self,
+ metric: str,
+ with_index: bool = True,
+ ) -> list:
+ if with_index:
+ return self._metrics[metric]
+ return [value[2] for value in self._metrics[metric]]
+
+ def as_json(
+ self,
+ output_dir: str = None,
+ filename: str = None
+ ) -> Dict[str, List[Dict[str, float]]]:
+ metrics_data = {}
+ for metric, values in self._metrics.items():
+ metrics_data[metric] = [
+ {'epoch': epoch, 'step': step, 'value': value}
+ for epoch, step, value
+ in values
+ ]
+
+ if output_dir and filename:
+ if not os.path.exists(output_dir):
+ os.makedirs(output_dir)
+ filepath = os.path.join(output_dir, filename)
+ with open(filepath, 'w') as file:
+ json.dump(metrics_data, file, indent=4)
+
+ return metrics_data
+
+
+def aggregate_by_epoch(metric_values: List[Tuple[int, int, float]]) -> Dict[int, float]:
+ """
+ Aggregates metric values by epoch, calculating the average for each epoch.
+
+ Args:
+ metric_values (List[Tuple[int, int, float]]): A list of tuples containing
+ (epoch, step, value) for each metric measurement.
+
+ Returns:
+ Dict[int, float]: A dictionary with epochs as keys and average metric values as values.
+ """
+ epoch_data = defaultdict(list)
+ for epoch, step, value in metric_values:
+ epoch_data[epoch].append(value)
+ avg_per_epoch = {
+ epoch: sum(values) / len(values)
+ for epoch, values
+ in epoch_data.items()
+ }
+ return avg_per_epoch
+
+
+def save_metric_plots(
+ training_tracker: MetricsTracker,
+ validation_tracker: MetricsTracker,
+ output_dir: str
+):
+ """
+ Saves plots of training and validation metrics over epochs.
+
+ Args:
+ training_tracker (MetricsTracker): Tracker containing training metrics.
+ validation_tracker (MetricsTracker): Tracker containing validation metrics.
+ output_dir (str): Directory to save the generated plots.
+
+ Returns:
+ None
+ """
+ if not os.path.exists(output_dir):
+ os.makedirs(output_dir)
+
+ training_metrics = training_tracker.describe_metrics()
+ validation_metrics = validation_tracker.describe_metrics()
+ all_metrics = set(training_metrics + validation_metrics)
+
+ for metric in all_metrics:
+ plt.figure(figsize=(8, 6))
+
+ if metric in training_metrics:
+ training_values = training_tracker.get_metric_values(
+ metric=metric, with_index=True)
+ training_avg_values = aggregate_by_epoch(training_values)
+ training_epochs = sorted(training_avg_values.keys())
+ training_vals = [training_avg_values[epoch] for epoch in training_epochs]
+ plt.plot(
+ training_epochs,
+ training_vals,
+ label=f'Training {metric}',
+ marker='o',
+ linestyle='-',
+ color='blue'
+ )
+
+ if metric in validation_metrics:
+ validation_values = validation_tracker.get_metric_values(
+ metric=metric, with_index=True)
+ validation_avg_values = aggregate_by_epoch(validation_values)
+ validation_epochs = sorted(validation_avg_values.keys())
+ validation_vals = [
+ validation_avg_values[epoch]
+ for epoch
+ in validation_epochs
+ ]
+ plt.plot(
+ validation_epochs,
+ validation_vals,
+ label=f'Validation {metric}',
+ marker='o',
+ linestyle='--',
+ color='orange'
+ )
+
+ plt.title(f'{metric.capitalize()} over Epochs')
+ plt.xlabel('Epoch')
+ plt.ylabel(f'{metric.capitalize()} Value')
+ plt.legend()
+ plt.grid(True)
+ plt.savefig(f'{output_dir}/{metric}_plot.png')
+ plt.close()
+
+
+
+def display_results(
+ prompts: List[str],
+ expected_responses: List[str],
+ generated_texts: List[str],
+ images: List[Image.Image]
+) -> None:
+ """
+ Display the results of model inference in IPython environments.
+
+ This function attempts to display the results (prompts, expected responses,
+ generated texts, and images) in an HTML format if running in an IPython
+ environment. If not in IPython or if there's an ImportError, it silently passes.
+
+ Args:
+ prompts (List[str]): List of input prompts.
+ expected_responses (List[str]): List of expected responses.
+ generated_texts (List[str]): List of texts generated by the model.
+ images (List[Image.Image]): List of input images.
+
+ Returns:
+ None
+ """
+ try:
+ import IPython
+ if IPython.get_ipython() is not None:
+ from IPython.display import display, HTML
+ html_out = create_html_output(prompts, expected_responses, generated_texts, images)
+ display(HTML(html_out))
+ except ImportError:
+ pass # Skip visualization if required libraries are not available
+
+
+def create_html_output(
+ prompts: List[str],
+ expected_responses: List[str],
+ generated_texts: List[str],
+ images: List[Image.Image]
+) -> str:
+ """
+ Create an HTML string to display the results of model inference.
+
+ This function generates an HTML string that includes styled divs for each
+ result, containing the input image, prompt, expected response, and generated text.
+
+ Args:
+ prompts (List[str]): List of input prompts.
+ expected_responses (List[str]): List of expected responses.
+ generated_texts (List[str]): List of texts generated by the model.
+ images (List[Image.Image]): List of input images.
+
+ Returns:
+ str: An HTML string containing the formatted results.
+ """
+ html_out = ""
+ count = min(8, len(images)) # Display up to 8 examples
+ for i in range(count):
+ html_out += f"""
+ 👉 read more
- -- -- **load image** - - ```python - import cv2 - - image = cv2.imread("...") - ``` - -- **create and refine marks** - - ```python - import maestro - - generator = maestro.SegmentAnythingMarkGenerator(device='cuda') - marks = generator.generate(image=image) - marks = maestro.refine_marks(marks=marks) - ``` - -- **visualize marks** - - ```python - mark_visualizer = maestro.MarkVisualizer() - marked_image = mark_visualizer.visualize(image=image, marks=marks) - ``` -  - -- **prompt** - - ```python - prompt = "Find dog." - - response = maestro.prompt_image(api_key=api_key, image=marked_image, prompt=prompt) - ``` - - ``` - >>> "The dog is prominently featured in the center of the image with the label [9]." - ``` - -- **extract related marks** - - ```python - masks = maestro.extract_relevant_masks(text=response, detections=refined_marks) - ``` - - ``` - >>> {'6': array([ - ... [False, False, False, ..., False, False, False], - ... [False, False, False, ..., False, False, False], - ... [False, False, False, ..., False, False, False], - ... ..., - ... [ True, True, True, ..., False, False, False], - ... [ True, True, True, ..., False, False, False], - ... [ True, True, True, ..., False, False, False]]) - ... } - ``` - -
+
+ """
+ return html_out
+
+
+def render_inline(image: Image.Image, resize: Tuple[int, int] = (256, 256)) -> str:
+ """
+ Convert an image into an inline HTML string.
+
+ This function takes an image, resizes it, and converts it to a base64-encoded
+ string that can be used as the source for an HTML img tag.
+
+ Args:
+ image (Image.Image): The input image to be converted.
+ resize (Tuple[int, int], optional): The dimensions to resize the image to.
+ Defaults to (256, 256).
+
+ Returns:
+ str: A string containing the data URI for the image, ready to be used
+ in an HTML img tag's src attribute.
+ """
+ image = image.resize(resize)
+ with io.BytesIO() as buffer:
+ image.save(buffer, format='jpeg')
+ image_b64 = base64.b64encode(buffer.getvalue()).decode()
+ return f"data:image/jpeg;base64,{image_b64}"
\ No newline at end of file
diff --git a/maestro/trainer/common/utils/reproducibility.py b/maestro/trainer/common/utils/reproducibility.py
new file mode 100644
index 0000000..69e6f1f
--- /dev/null
+++ b/maestro/trainer/common/utils/reproducibility.py
@@ -0,0 +1,24 @@
+import os
+import random
+from typing import Optional
+
+import torch
+import numpy as np
+
+from maestro.trainer.common.configuration.env import SEED_ENV, DEFAULT_SEED
+
+
+def make_it_reproducible(
+ seed: Optional[int] = None,
+ disable_cudnn_benchmark: bool = True,
+ avoid_non_deterministic_algorithms: bool = True,
+) -> None:
+ if seed is None:
+ seed = int(os.getenv(SEED_ENV, DEFAULT_SEED))
+ random.seed(seed)
+ torch.manual_seed(seed)
+ np.random.seed(seed)
+ if avoid_non_deterministic_algorithms:
+ torch.use_deterministic_algorithms(True)
+ if disable_cudnn_benchmark:
+ torch.backends.cudnn.benchmark = False
diff --git a/maestro/trainer/models/__init__.py b/maestro/trainer/models/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/maestro/trainer/models/florence_2/__init__.py b/maestro/trainer/models/florence_2/__init__.py
new file mode 100644
index 0000000..77d45b3
--- /dev/null
+++ b/maestro/trainer/models/florence_2/__init__.py
@@ -0,0 +1 @@
+from maestro.trainer.models.florence_2.core import TrainingConfiguration, train
diff --git a/maestro/trainer/models/florence_2/checkpoints.py b/maestro/trainer/models/florence_2/checkpoints.py
new file mode 100644
index 0000000..a489f2b
--- /dev/null
+++ b/maestro/trainer/models/florence_2/checkpoints.py
@@ -0,0 +1,119 @@
+import os
+from typing import Optional, Tuple
+
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+from maestro.trainer.common.configuration.env import CUDA_DEVICE_ENV, \
+ DEFAULT_CUDA_DEVICE
+
+DEFAULT_FLORENCE2_MODEL_ID = "microsoft/Florence-2-base-ft"
+DEFAULT_FLORENCE2_MODEL_REVISION = "refs/pr/20"
+DEVICE = torch.device("cpu") \
+ if not torch.cuda.is_available() \
+ else os.getenv(CUDA_DEVICE_ENV, DEFAULT_CUDA_DEVICE)
+
+
+class CheckpointManager:
+ """Manages checkpoints for model training.
+
+ This class handles saving and retrieving model checkpoints during training.
+
+ Attributes:
+ training_dir (str): Directory where checkpoints will be saved.
+ best_val_loss (float): Best validation loss achieved so far.
+ latest_checkpoint_dir (str): Directory for the latest checkpoint.
+ best_checkpoint_dir (str): Directory for the best checkpoint.
+ """
+
+ def __init__(self, training_dir: str):
+ """Initializes the CheckpointManager.
+
+ Args:
+ training_dir (str): Directory where checkpoints will be saved.
+ """
+ self.training_dir = training_dir
+ self.best_val_loss = float('inf')
+ self.latest_checkpoint_dir = os.path.join(training_dir, "checkpoints", "latest")
+ self.best_checkpoint_dir = os.path.join(training_dir, "checkpoints", "best")
+
+ def save_latest(self, processor: AutoProcessor, model: AutoModelForCausalLM):
+ """Saves the latest model checkpoint.
+
+ Args:
+ processor (AutoProcessor): The processor to save.
+ model (AutoModelForCausalLM): The model to save.
+ """
+ save_model(self.latest_checkpoint_dir, processor, model)
+
+ def save_best(self, processor: AutoProcessor, model: AutoModelForCausalLM, val_loss: float):
+ """Saves the best model checkpoint if the validation loss improves.
+
+ Args:
+ processor (AutoProcessor): The processor to save.
+ model (AutoModelForCausalLM): The model to save.
+ val_loss (float): The current validation loss.
+ """
+ if val_loss < self.best_val_loss:
+ self.best_val_loss = val_loss
+ save_model(self.best_checkpoint_dir, processor, model)
+ print(f"New best model saved with validation loss: {self.best_val_loss}")
+
+ def get_best_model_path(self):
+ """Returns the path to the best model checkpoint.
+
+ Returns:
+ str: Path to the best model checkpoint.
+ """
+ return self.best_checkpoint_dir
+
+
+def save_model(
+ target_dir: str,
+ processor: AutoProcessor,
+ model: AutoModelForCausalLM,
+) -> None:
+ """Saves the model and processor to the specified directory.
+
+ Args:
+ target_dir (str): Directory where the model and processor will be saved.
+ processor (AutoProcessor): The processor to save.
+ model (AutoModelForCausalLM): The model to save.
+ """
+ os.makedirs(target_dir, exist_ok=True)
+ processor.save_pretrained(target_dir)
+ model.save_pretrained(target_dir)
+
+
+def load_model(
+ model_id_or_path: str = DEFAULT_FLORENCE2_MODEL_ID,
+ revision: str = DEFAULT_FLORENCE2_MODEL_REVISION,
+ device: torch.device = DEVICE,
+ cache_dir: Optional[str] = None,
+) -> Tuple[AutoProcessor, AutoModelForCausalLM]:
+ """Loads a Florence-2 model and its associated processor.
+
+ Args:
+ model_id_or_path: The identifier or path of the model to load.
+ revision: The specific model revision to use.
+ device: The device to load the model onto.
+ cache_dir: Directory to cache the downloaded model files.
+
+ Returns:
+ A tuple containing the loaded processor and model.
+
+ Raises:
+ ValueError: If the model or processor cannot be loaded.
+ """
+ processor = AutoProcessor.from_pretrained(
+ model_id_or_path,
+ trust_remote_code=True,
+ revision=revision,
+ )
+ model = AutoModelForCausalLM.from_pretrained(
+ model_id_or_path,
+ trust_remote_code=True,
+ revision=revision,
+ cache_dir=cache_dir,
+ ).to(device)
+ return processor, model
diff --git a/maestro/trainer/models/florence_2/core.py b/maestro/trainer/models/florence_2/core.py
new file mode 100644
index 0000000..8ef20ce
--- /dev/null
+++ b/maestro/trainer/models/florence_2/core.py
@@ -0,0 +1,423 @@
+import os
+from dataclasses import dataclass, field, replace
+from typing import List, Literal, Optional, Tuple, Union
+
+import torch
+from peft import LoraConfig, PeftModel, get_peft_model
+from torch.optim import Adam, AdamW, Optimizer, SGD
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, AutoProcessor, get_scheduler
+
+from maestro.trainer.common.utils.file_system import create_new_run_directory
+from maestro.trainer.common.utils.metrics import (
+ BaseMetric,
+ MetricsTracker,
+ display_results,
+ save_metric_plots,
+ MeanAveragePrecisionMetric
+)
+from maestro.trainer.common.utils.reproducibility import make_it_reproducible
+from maestro.trainer.models.florence_2.checkpoints import (
+ CheckpointManager,
+ load_model,
+ DEFAULT_FLORENCE2_MODEL_ID,
+ DEFAULT_FLORENCE2_MODEL_REVISION,
+ DEVICE
+)
+from maestro.trainer.models.florence_2.data_loading import prepare_data_loaders
+from maestro.trainer.models.florence_2.metrics import (
+ extract_unique_detection_dataset_classes,
+ postprocess_florence2_output_for_mean_average_precision,
+ run_predictions,
+)
+from maestro.trainer.models.paligemma.training import LoraInitLiteral
+
+
+@dataclass(frozen=True)
+class TrainingConfiguration:
+ """Configuration for training a Florence-2 model.
+
+ This class encapsulates all the parameters needed for training a Florence-2 model,
+ including dataset paths, model specifications, training hyperparameters, and output
+ settings.
+
+ Attributes:
+ dataset (str): Path to the dataset used for training.
+ model_id (str): Identifier for the Florence-2 model.
+ revision (str): Revision of the model to use.
+ device (torch.device): Device to use for training.
+ cache_dir (Optional[str]): Directory to cache the model.
+ epochs (int): Number of training epochs.
+ optimizer (Literal["sgd", "adamw", "adam"]): Optimizer to use for training.
+ lr (float): Learning rate for the optimizer.
+ lr_scheduler (Literal["linear", "cosine", "polynomial"]): Learning rate
+ scheduler.
+ batch_size (int): Batch size for training.
+ val_batch_size (Optional[int]): Batch size for validation.
+ num_workers (int): Number of workers for data loading.
+ val_num_workers (Optional[int]): Number of workers for validation data loading.
+ lora_r (int): Rank of the LoRA update matrices.
+ lora_alpha (int): Scaling factor for the LoRA update.
+ lora_dropout (float): Dropout probability for LoRA layers.
+ bias (Literal["none", "all", "lora_only"]): Which bias to train.
+ use_rslora (bool): Whether to use RSLoRA.
+ init_lora_weights (Union[bool, LoraInitLiteral]): How to initialize LoRA
+ weights.
+ output_dir (str): Directory to save output files.
+ metrics (List[BaseMetric]): List of metrics to track during training.
+ """
+ dataset: str
+ model_id: str = DEFAULT_FLORENCE2_MODEL_ID
+ revision: str = DEFAULT_FLORENCE2_MODEL_REVISION
+ device: torch.device = DEVICE
+ cache_dir: Optional[str] = None
+ epochs: int = 10
+ optimizer: Literal["sgd", "adamw", "adam"] = "adamw"
+ lr: float = 1e-5
+ lr_scheduler: Literal["linear", "cosine", "polynomial"] = "linear"
+ batch_size: int = 4
+ val_batch_size: Optional[int] = None
+ num_workers: int = 0
+ val_num_workers: Optional[int] = None
+ lora_r: int = 8
+ lora_alpha: int = 8
+ lora_dropout: float = 0.05
+ bias: Literal["none", "all", "lora_only"] = "none"
+ use_rslora: bool = True
+ init_lora_weights: Union[bool, LoraInitLiteral] = "gaussian"
+ output_dir: str = "./training/florence-2"
+ metrics: List[BaseMetric] = field(default_factory=list)
+
+
+def train(config: TrainingConfiguration) -> None:
+ make_it_reproducible(avoid_non_deterministic_algorithms=False)
+ run_dir = create_new_run_directory(
+ base_output_dir=config.output_dir,
+ )
+ config = replace(
+ config,
+ output_dir=run_dir,
+ )
+ checkpoint_manager = CheckpointManager(run_dir)
+
+ processor, model = load_model(
+ model_id_or_path=config.model_id,
+ revision=config.revision,
+ device=config.device,
+ cache_dir=config.cache_dir,
+ )
+ train_loader, val_loader, test_loader = prepare_data_loaders(
+ dataset_location=config.dataset,
+ train_batch_size=config.batch_size,
+ processor=processor,
+ device=config.device,
+ num_workers=config.num_workers,
+ test_loaders_workers=config.val_num_workers,
+ )
+ peft_model = prepare_peft_model(
+ model=model,
+ r=config.lora_r,
+ lora_alpha=config.lora_alpha,
+ lora_dropout=config.lora_dropout,
+ bias=config.bias,
+ use_rslora=config.use_rslora,
+ init_lora_weights=config.init_lora_weights,
+ revision=config.revision,
+ )
+ training_metrics_tracker = MetricsTracker.init(metrics=["loss"])
+ metrics = ["loss"]
+ for metric in config.metrics:
+ metrics += metric.describe()
+ validation_metrics_tracker = MetricsTracker.init(metrics=metrics)
+
+ run_training_loop(
+ processor=processor,
+ model=peft_model,
+ data_loaders=(train_loader, val_loader),
+ config=config,
+ training_metrics_tracker=training_metrics_tracker,
+ validation_metrics_tracker=validation_metrics_tracker,
+ checkpoint_manager=checkpoint_manager
+ )
+
+ save_metric_plots(
+ training_tracker=training_metrics_tracker,
+ validation_tracker=validation_metrics_tracker,
+ output_dir=os.path.join(config.output_dir, "metrics"),
+ )
+ training_metrics_tracker.as_json(
+ output_dir=os.path.join(config.output_dir, "metrics"),
+ filename="training.json")
+ validation_metrics_tracker.as_json(
+ output_dir=os.path.join(config.output_dir, "metrics"),
+ filename="validation.json")
+
+ # Log out paths for latest and best checkpoints
+ print(f"Latest checkpoint saved at: {checkpoint_manager.latest_checkpoint_dir}")
+ print(f"Best checkpoint saved at: {checkpoint_manager.best_checkpoint_dir}")
+
+
+def prepare_peft_model(
+ model: AutoModelForCausalLM,
+ r: int = 8,
+ lora_alpha: int = 8,
+ lora_dropout: float = 0.05,
+ bias: Literal["none", "all", "lora_only"] = "none",
+ inference_mode: bool = False,
+ use_rslora: bool = True,
+ init_lora_weights: Union[bool, LoraInitLiteral] = "gaussian",
+ revision: str = DEFAULT_FLORENCE2_MODEL_REVISION,
+) -> PeftModel:
+ config = LoraConfig(
+ r=r,
+ lora_alpha=lora_alpha,
+ target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "linear", "Conv2d", "lm_head", "fc2"],
+ task_type="CAUSAL_LM",
+ lora_dropout=lora_dropout,
+ bias=bias,
+ inference_mode=inference_mode,
+ use_rslora=use_rslora,
+ init_lora_weights=init_lora_weights,
+ revision=revision,
+ )
+ peft_model = get_peft_model(model, config)
+ peft_model.print_trainable_parameters()
+ return peft_model.to(model.device)
+
+
+def run_training_loop(
+ processor: AutoProcessor,
+ model: PeftModel,
+ data_loaders: Tuple[DataLoader, Optional[DataLoader]],
+ config: TrainingConfiguration,
+ training_metrics_tracker: MetricsTracker,
+ validation_metrics_tracker: MetricsTracker,
+ checkpoint_manager: CheckpointManager,
+) -> None:
+ train_loader, val_loader = data_loaders
+ optimizer = get_optimizer(model=model, config=config)
+ total_steps = config.epochs * len(train_loader)
+ lr_scheduler = get_scheduler(
+ name=config.lr_scheduler,
+ optimizer=optimizer,
+ num_warmup_steps=0,
+ num_training_steps=total_steps,
+ )
+ for epoch in range(config.epochs):
+ run_training_epoch(
+ processor=processor,
+ model=model,
+ train_loader=train_loader,
+ val_loader=val_loader,
+ epoch=epoch + 1,
+ config=config,
+ optimizer=optimizer,
+ lr_scheduler=lr_scheduler,
+ training_metrics_tracker=training_metrics_tracker,
+ validation_metrics_tracker=validation_metrics_tracker,
+ checkpoint_manager=checkpoint_manager
+ )
+
+
+def run_training_epoch(
+ processor: AutoProcessor,
+ model: PeftModel,
+ train_loader: DataLoader,
+ val_loader: Optional[DataLoader],
+ epoch: int,
+ config: TrainingConfiguration,
+ optimizer: Optimizer,
+ lr_scheduler: LRScheduler,
+ training_metrics_tracker: MetricsTracker,
+ validation_metrics_tracker: MetricsTracker,
+ checkpoint_manager: CheckpointManager,
+) -> None:
+ model.train()
+ training_losses: List[float] = []
+
+ with tqdm(total=len(train_loader), desc=f"Epoch {epoch}/{config.epochs}", unit="batch") as pbar:
+ for step_id, (inputs, answers) in enumerate(train_loader):
+ input_ids = inputs["input_ids"]
+ pixel_values = inputs["pixel_values"]
+ labels = processor.tokenizer(
+ text=answers,
+ return_tensors="pt",
+ padding=True,
+ return_token_type_ids=False
+ ).input_ids.to(config.device)
+ outputs = model(input_ids=input_ids, pixel_values=pixel_values, labels=labels)
+ loss = outputs.loss
+ loss.backward()
+ optimizer.step()
+ lr_scheduler.step()
+ optimizer.zero_grad()
+ loss = loss.item()
+ training_metrics_tracker.register(
+ metric="loss",
+ epoch=epoch,
+ step=step_id + 1,
+ value=loss,
+ )
+ training_losses.append(loss)
+
+ # Update progress bar
+ last_100_losses = training_losses[-100:]
+ loss_moving_average = sum(last_100_losses) / len(last_100_losses) if last_100_losses else 0.0
+ pbar.set_postfix({"Loss": f"{loss_moving_average:.4f}"})
+ pbar.update(1)
+
+ # Save checkpoints based on training loss if no validation loader
+ if val_loader is None or len(val_loader) == 0:
+ train_loss = sum(training_losses) / len(training_losses)
+ checkpoint_manager.save_latest(processor, model)
+ checkpoint_manager.save_best(processor, model, train_loss)
+ return
+
+ run_validation_epoch(
+ processor=processor,
+ model=model,
+ loader=val_loader,
+ epoch_number=epoch,
+ config=config,
+ metrics_tracker=validation_metrics_tracker,
+ )
+
+ val_loss = validation_metrics_tracker.get_metric_values("loss")[-1][2]
+ checkpoint_manager.save_latest(processor, model)
+ checkpoint_manager.save_best(processor, model, val_loss)
+
+
+def run_validation_epoch(
+ processor: AutoProcessor,
+ model: Union[PeftModel, AutoModelForCausalLM],
+ loader: DataLoader,
+ config: TrainingConfiguration,
+ metrics_tracker: MetricsTracker,
+ epoch_number: int
+) -> None:
+ val_loss = 0.0
+ with torch.no_grad():
+ for inputs, targets in loader:
+ input_ids = inputs["input_ids"]
+ pixel_values = inputs["pixel_values"]
+ labels = processor.tokenizer(
+ text=targets,
+ return_tensors="pt",
+ padding=True,
+ return_token_type_ids=False
+ ).input_ids.to(config.device)
+ outputs = model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ labels=labels
+ )
+ loss = outputs.loss
+ val_loss += loss.item()
+ avg_val_loss = val_loss / len(loader)
+ metrics_tracker.register(
+ metric="loss",
+ epoch=epoch_number,
+ step=1,
+ value=avg_val_loss,
+ )
+ # Run inference once for all metrics
+ prompts, expected_responses, generated_texts, images = run_predictions(
+ dataset=loader.dataset,
+ processor=processor,
+ model=model,
+ device=config.device,
+ )
+
+ metrics_results = {"loss": avg_val_loss}
+
+ for metric in config.metrics:
+ if isinstance(metric, MeanAveragePrecisionMetric):
+ classes = extract_unique_detection_dataset_classes(loader.dataset)
+ targets, predictions = postprocess_florence2_output_for_mean_average_precision(
+ expected_responses=expected_responses,
+ generated_texts=generated_texts,
+ images=images,
+ classes=classes,
+ processor=processor
+ )
+ result = metric.compute(targets=targets, predictions=predictions)
+ for key, value in result.items():
+ metrics_tracker.register(
+ metric=key,
+ epoch=epoch_number,
+ step=1,
+ value=value,
+ )
+ metrics_results[key] = value
+
+ print("Validation Metrics:", ", ".join([f"{k}: {v:.4f}" for k, v in metrics_results.items()]))
+
+ # Display inference results in IPython environments
+ display_results(prompts, expected_responses, generated_texts, images)
+
+
+def get_optimizer(model: PeftModel, config: TrainingConfiguration) -> Optimizer:
+ optimizer_type = config.optimizer.lower()
+ if optimizer_type == "adamw":
+ return AdamW(model.parameters(), lr=config.lr)
+ if optimizer_type == "adam":
+ return Adam(model.parameters(), lr=config.lr)
+ if optimizer_type == "sgd":
+ return SGD(model.parameters(), lr=config.lr)
+ raise ValueError(f"Unsupported optimizer: {config.optimizer}")
+
+
+def evaluate(config: TrainingConfiguration) -> None:
+ processor, model = load_model(
+ model_id_or_path=config.model_id,
+ revision=config.revision,
+ device=config.device,
+ cache_dir=config.cache_dir,
+ )
+ train_loader, val_loader, test_loader = prepare_data_loaders(
+ dataset_location=config.dataset,
+ train_batch_size=config.batch_size,
+ processor=processor,
+ device=config.device,
+ num_workers=config.num_workers,
+ test_loaders_workers=config.val_num_workers,
+ )
+ evaluation_loader = test_loader if test_loader is not None else val_loader
+
+ metrics = []
+ for metric in config.metrics:
+ metrics += metric.describe()
+ evaluation_metrics_tracker = MetricsTracker.init(metrics=metrics)
+
+ # Run inference once for all metrics
+ _, expected_responses, generated_texts, images = run_predictions(
+ dataset=evaluation_loader.dataset,
+ processor=processor,
+ model=model,
+ device=config.device,
+ )
+
+ for metric in config.metrics:
+ if isinstance(metric, MeanAveragePrecisionMetric):
+ classes = extract_unique_detection_dataset_classes(train_loader.dataset)
+ targets, predictions = postprocess_florence2_output_for_mean_average_precision(
+ expected_responses=expected_responses,
+ generated_texts=generated_texts,
+ images=images,
+ classes=classes,
+ processor=processor
+ )
+ result = metric.compute(targets=targets, predictions=predictions)
+ for key, value in result.items():
+ evaluation_metrics_tracker.register(
+ metric=key,
+ epoch=1,
+ step=1,
+ value=value,
+ )
+
+ evaluation_metrics_tracker.as_json(
+ output_dir=os.path.join(config.output_dir, "metrics"),
+ filename="evaluation.json")
diff --git a/maestro/trainer/models/florence_2/data_loading.py b/maestro/trainer/models/florence_2/data_loading.py
new file mode 100644
index 0000000..4873003
--- /dev/null
+++ b/maestro/trainer/models/florence_2/data_loading.py
@@ -0,0 +1,110 @@
+import logging
+import os
+from functools import partial
+from typing import Optional, Tuple, List
+
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from transformers import AutoProcessor
+
+from maestro.trainer.common.data_loaders.datasets import DetectionDataset
+
+
+def prepare_data_loaders(
+ dataset_location: str,
+ train_batch_size: int,
+ processor: AutoProcessor,
+ device: torch.device,
+ num_workers: int = 0,
+ test_batch_size: Optional[int] = None,
+ test_loaders_workers: Optional[int] = None,
+) -> Tuple[
+ DataLoader,
+ Optional[DataLoader],
+ Optional[DataLoader],
+]:
+ test_batch_size = test_batch_size or train_batch_size
+ test_loaders_workers = test_loaders_workers or num_workers
+ train_data_loader = prepare_detection_data_loader(
+ dataset_location=dataset_location,
+ split_name="train",
+ batch_size=train_batch_size,
+ processor=processor,
+ device=device,
+ num_workers=num_workers,
+ shuffle=True,
+ )
+ if train_data_loader is None:
+ raise RuntimeError("Could not initialise train data loader")
+ valid_data_loader = prepare_detection_data_loader(
+ dataset_location=dataset_location,
+ split_name="valid",
+ batch_size=test_batch_size,
+ processor=processor,
+ device=device,
+ num_workers=test_loaders_workers,
+ shuffle=False,
+ )
+ test_data_loader = prepare_detection_data_loader(
+ dataset_location=dataset_location,
+ split_name="test",
+ batch_size=test_batch_size,
+ processor=processor,
+ device=device,
+ num_workers=test_loaders_workers,
+ shuffle=False,
+ )
+ return train_data_loader, valid_data_loader, test_data_loader
+
+
+def prepare_detection_data_loader(
+ dataset_location: str,
+ split_name: str,
+ batch_size: int,
+ processor: AutoProcessor,
+ device: torch.device,
+ num_workers: int = 0,
+ shuffle: bool = True,
+) -> Optional[DataLoader]:
+ dataset = prepare_detection_dataset(
+ dataset_location=dataset_location,
+ split_name=split_name,
+ )
+ if dataset is None:
+ return None
+ return DataLoader(
+ dataset,
+ batch_size=batch_size,
+ collate_fn=partial(collate_fn, processor=processor, device=device),
+ num_workers=num_workers,
+ shuffle=shuffle,
+ )
+
+
+def prepare_detection_dataset(
+ dataset_location: str,
+ split_name: str,
+) -> Optional[DetectionDataset]:
+ image_directory_path = os.path.join(dataset_location, split_name)
+ jsonl_file_path = os.path.join(dataset_location, split_name, "annotations.jsonl")
+ if not os.path.exists(image_directory_path):
+ logging.warning(f"Could not data directory: {image_directory_path}")
+ return None
+ if not os.path.exists(jsonl_file_path):
+ logging.warning(f"Could not find JSONL file: {jsonl_file_path}")
+ return None
+ return DetectionDataset(
+ jsonl_file_path=jsonl_file_path,
+ image_directory_path=image_directory_path,
+ )
+
+
+def collate_fn(
+ batch: Tuple[List[str], List[str], List[Image.Image]],
+ processor: AutoProcessor,
+ device: torch.device,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ questions, answers, images = zip(*batch)
+ inputs = processor(text=list(questions), images=list(images), return_tensors="pt", padding=True).to(device)
+ return inputs, answers
diff --git a/maestro/trainer/models/florence_2/entrypoint.py b/maestro/trainer/models/florence_2/entrypoint.py
new file mode 100644
index 0000000..f18d81e
--- /dev/null
+++ b/maestro/trainer/models/florence_2/entrypoint.py
@@ -0,0 +1,219 @@
+import dataclasses
+from typing import Optional, Annotated, List, Dict, Type
+
+import rich
+import torch
+import typer
+
+from maestro.trainer.models.florence_2.checkpoints import DEFAULT_FLORENCE2_MODEL_ID, \
+ DEFAULT_FLORENCE2_MODEL_REVISION, DEVICE
+from maestro.trainer.models.florence_2.core import TrainingConfiguration
+from maestro.trainer.models.florence_2.core import train as florence2_train
+from maestro.trainer.models.florence_2.core import evaluate as florence2_evaluate
+from maestro.trainer.common.utils.metrics import BaseMetric, MeanAveragePrecisionMetric
+
+florence_2_app = typer.Typer(help="Fine-tune and evaluate Florence 2 model")
+
+
+METRIC_CLASSES: Dict[str, Type[BaseMetric]] = {
+ "mean_average_precision": MeanAveragePrecisionMetric,
+}
+
+
+def parse_metrics(metrics: List[str]) -> List[BaseMetric]:
+ metric_objects = []
+ for metric_name in metrics:
+ metric_class = METRIC_CLASSES.get(metric_name.lower())
+ if metric_class:
+ metric_objects.append(metric_class())
+ else:
+ raise ValueError(f"Unsupported metric: {metric_name}")
+ return metric_objects
+
+
+@florence_2_app.command(
+ help="Train Florence 2 model",
+ context_settings={"allow_extra_args": True, "ignore_unknown_options": True}
+)
+def train(
+ dataset: Annotated[
+ str,
+ typer.Option("--dataset", help="Path to the dataset used for training"),
+ ],
+ model_id: Annotated[
+ str,
+ typer.Option("--model_id", help="Identifier for the Florence-2 model"),
+ ] = DEFAULT_FLORENCE2_MODEL_ID,
+ revision: Annotated[
+ str,
+ typer.Option("--revision", help="Revision of the model to use"),
+ ] = DEFAULT_FLORENCE2_MODEL_REVISION,
+ device: Annotated[
+ str,
+ typer.Option("--device", help="Device to use for training"),
+ ] = DEVICE,
+ cache_dir: Annotated[
+ Optional[str],
+ typer.Option("--cache_dir", help="Directory to cache the model"),
+ ] = None,
+ epochs: Annotated[
+ int,
+ typer.Option("--epochs", help="Number of training epochs"),
+ ] = 10,
+ optimizer: Annotated[
+ str,
+ typer.Option("--optimizer", help="Optimizer to use for training"),
+ ] = "adamw",
+ lr: Annotated[
+ float,
+ typer.Option("--lr", help="Learning rate for the optimizer"),
+ ] = 1e-5,
+ lr_scheduler: Annotated[
+ str,
+ typer.Option("--lr_scheduler", help="Learning rate scheduler"),
+ ] = "linear",
+ batch_size: Annotated[
+ int,
+ typer.Option("--batch_size", help="Batch size for training"),
+ ] = 4,
+ val_batch_size: Annotated[
+ Optional[int],
+ typer.Option("--val_batch_size", help="Batch size for validation"),
+ ] = None,
+ num_workers: Annotated[
+ int,
+ typer.Option("--num_workers", help="Number of workers for data loading"),
+ ] = 0,
+ val_num_workers: Annotated[
+ Optional[int],
+ typer.Option("--val_num_workers", help="Number of workers for validation data loading"),
+ ] = None,
+ lora_r: Annotated[
+ int,
+ typer.Option("--lora_r", help="Rank of the LoRA update matrices"),
+ ] = 8,
+ lora_alpha: Annotated[
+ int,
+ typer.Option("--lora_alpha", help="Scaling factor for the LoRA update"),
+ ] = 8,
+ lora_dropout: Annotated[
+ float,
+ typer.Option("--lora_dropout", help="Dropout probability for LoRA layers"),
+ ] = 0.05,
+ bias: Annotated[
+ str,
+ typer.Option("--bias", help="Which bias to train"),
+ ] = "none",
+ use_rslora: Annotated[
+ bool,
+ typer.Option("--use_rslora/--no_use_rslora", help="Whether to use RSLoRA"),
+ ] = True,
+ init_lora_weights: Annotated[
+ str,
+ typer.Option("--init_lora_weights", help="How to initialize LoRA weights"),
+ ] = "gaussian",
+ output_dir: Annotated[
+ str,
+ typer.Option("--output_dir", help="Directory to save output files"),
+ ] = "./training/florence-2",
+ metrics: Annotated[
+ List[str],
+ typer.Option("--metrics", help="List of metrics to track during training"),
+ ] = [],
+) -> None:
+ metric_objects = parse_metrics(metrics)
+ config = TrainingConfiguration(
+ dataset=dataset,
+ model_id=model_id,
+ revision=revision,
+ device=torch.device(device),
+ cache_dir=cache_dir,
+ epochs=epochs,
+ optimizer=optimizer,
+ lr=lr,
+ lr_scheduler=lr_scheduler,
+ batch_size=batch_size,
+ val_batch_size=val_batch_size,
+ num_workers=num_workers,
+ val_num_workers=val_num_workers,
+ lora_r=lora_r,
+ lora_alpha=lora_alpha,
+ lora_dropout=lora_dropout,
+ bias=bias,
+ use_rslora=use_rslora,
+ init_lora_weights=init_lora_weights,
+ output_dir=output_dir,
+ metrics=metric_objects
+ )
+ typer.echo(typer.style(
+ text="Training configuration",
+ fg=typer.colors.BRIGHT_GREEN,
+ bold=True
+ ))
+ rich.print(dataclasses.asdict(config))
+ florence2_train(config=config)
+
+
+@florence_2_app.command(help="Evaluate Florence 2 model")
+def evaluate(
+ dataset: Annotated[
+ str,
+ typer.Option("--dataset", help="Path to the dataset used for evaluation"),
+ ],
+ model_id: Annotated[
+ str,
+ typer.Option("--model_id", help="Identifier for the Florence-2 model"),
+ ] = DEFAULT_FLORENCE2_MODEL_ID,
+ revision: Annotated[
+ str,
+ typer.Option("--revision", help="Revision of the model to use"),
+ ] = DEFAULT_FLORENCE2_MODEL_REVISION,
+ device: Annotated[
+ str,
+ typer.Option("--device", help="Device to use for evaluation"),
+ ] = DEVICE,
+ cache_dir: Annotated[
+ Optional[str],
+ typer.Option("--cache_dir", help="Directory to cache the model"),
+ ] = None,
+ batch_size: Annotated[
+ int,
+ typer.Option("--batch_size", help="Batch size for evaluation"),
+ ] = 4,
+ num_workers: Annotated[
+ int,
+ typer.Option("--num_workers", help="Number of workers for data loading"),
+ ] = 0,
+ val_num_workers: Annotated[
+ Optional[int],
+ typer.Option("--val_num_workers", help="Number of workers for validation data loading"),
+ ] = None,
+ output_dir: Annotated[
+ str,
+ typer.Option("--output_dir", help="Directory to save output files"),
+ ] = "./evaluation/florence-2",
+ metrics: Annotated[
+ List[str],
+ typer.Option("--metrics", help="List of metrics to track during evaluation"),
+ ] = [],
+) -> None:
+ metric_objects = parse_metrics(metrics)
+ config = TrainingConfiguration(
+ dataset=dataset,
+ model_id=model_id,
+ revision=revision,
+ device=torch.device(device),
+ cache_dir=cache_dir,
+ batch_size=batch_size,
+ num_workers=num_workers,
+ val_num_workers=val_num_workers,
+ output_dir=output_dir,
+ metrics=metric_objects
+ )
+ typer.echo(typer.style(
+ text="Evaluation configuration",
+ fg=typer.colors.BRIGHT_GREEN,
+ bold=True
+ ))
+ rich.print(dataclasses.asdict(config))
+ florence2_evaluate(config=config)
diff --git a/maestro/trainer/models/florence_2/metrics.py b/maestro/trainer/models/florence_2/metrics.py
new file mode 100644
index 0000000..09e02bd
--- /dev/null
+++ b/maestro/trainer/models/florence_2/metrics.py
@@ -0,0 +1,83 @@
+import re
+from typing import List
+from typing import Tuple
+
+import numpy as np
+import supervision as sv
+import torch
+from PIL import Image
+from tqdm import tqdm
+from transformers import AutoProcessor, AutoModelForCausalLM
+
+from maestro.trainer.common.data_loaders.datasets import DetectionDataset
+
+DETECTION_CLASS_PATTERN = r"([a-zA-Z0-9 -]+)
+ })
+
+
+
+ Prompt: {html.escape(prompts[i])}
+ Expected: {html.escape(expected_responses[i])}
+ Generated: {html.escape(generated_texts[i])}
+