"
+) -> None:
+ """
+ Send a nicely formatted email with the list of downloadable links.
+
+ :param recipient: Email address to send the links to
+ :param links: List of presigned URLs
+ :param from_email: The email address sending this email
+ """
+ ses_client = boto3.client("ses")
+
+ links = create_presigned_urls(bucket_name=bucket_name, prefix=prefix)
+
+ # Email body
+ body_html = """
+
+
+
+ 🧠 Your Download Links from Geniusrise
+ We've prepared the files you requested. Below are the links to download them:
+
+ """
+ for link in links:
+ body_html += f"⬇️ Download Link
+ Please note that these links are valid for 24 hours only .
+ Thank you for using Geniusrise !
+
+
+ """
+
+ # Sending the email
+ try:
+ response = ses_client.send_email(
+ Source=from_email,
+ Destination={"ToAddresses": [recipient]},
+ Message={
+ "Subject": {"Data": "🧠 Your Download Links from Geniusrise"},
+ "Body": {"Html": {"Data": body_html}},
+ },
+ )
+ print(f"Email sent! Message ID: {response['MessageId']}")
+ except Exception as e:
+ print(f"An error occurred: {e}")
+
+
+def send_fine_tuning_email(
+ recipient: str, bucket_name: str, prefix: str, from_email: str = "Geniusrise "
+) -> None:
+ """
+ Send a nicely formatted email with the list of downloadable links.
+
+ :param recipient: Email address to send the links to
+ :param links: List of presigned URLs
+ :param from_email: The email address sending this email
+ """
+ ses_client = boto3.client("ses")
+
+ links = create_presigned_urls(bucket_name=bucket_name, prefix=prefix)
+
+ # Email body
+ body_html = """
+
+
+
+ 🧠 Your Fine-Tuned Model Download Links from Geniusrise
+ We've prepared the models you requested. Below are the links to download them:
+
+ """
+ for link in links:
+ body_html += f"⬇️ Download Link
+ Please note that these links are valid for 24 hours only .
+ Thank you for using Geniusrise !
+
+
+ """
+
+ # Sending the email
+ try:
+ response = ses_client.send_email(
+ Source=from_email,
+ Destination={"ToAddresses": [recipient]},
+ Message={
+ "Subject": {"Data": "🧠 Your Fine-Tuned Model Download Links from Geniusrise"},
+ "Body": {"Html": {"Data": body_html}},
+ },
+ )
+ print(f"Email sent! Message ID: {response['MessageId']}")
+ except Exception as e:
+ print(f"An error occurred: {e}")
diff --git a/geniusrise_text/base/fine_tune.py b/geniusrise_text/base/fine_tune.py
new file mode 100644
index 0000000..861437a
--- /dev/null
+++ b/geniusrise_text/base/fine_tune.py
@@ -0,0 +1,619 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from abc import abstractmethod
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+from accelerate import infer_auto_device_map, init_empty_weights
+from datasets import Dataset, DatasetDict
+from geniusrise import BatchInput, BatchOutput, Bolt, State
+from geniusrise.logging import setup_logger
+from peft import LoraConfig, get_peft_model
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+from transformers import AutoConfig, EvalPrediction, Trainer, TrainingArguments
+from trl import SFTTrainer
+
+from geniusrise_text.base.communication import send_fine_tuning_email
+from geniusrise_text.base.util import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING
+
+
+class TextFineTuner(Bolt):
+ """
+ A bolt for fine-tuning Hugging Face models.
+
+ This bolt uses the Hugging Face Transformers library to fine-tune a pre-trained model.
+ It uses the `Trainer` class from the Transformers library to handle the training.
+ """
+
+ def __init__(
+ self,
+ input: BatchInput,
+ output: BatchOutput,
+ state: State,
+ **kwargs,
+ ) -> None:
+ """
+ Initialize the bolt.
+
+ Args:
+ input (BatchInput): The batch input data.
+ output (BatchOutput): The output data.
+ state (State): The state manager.
+ evaluate (bool, optional): Whether to evaluate the model. Defaults to False.
+ **kwargs: Additional keyword arguments.
+ """
+ super().__init__(input=input, output=output, state=state)
+ self.input = input
+ self.output = output
+ self.state = state
+
+ self.log = setup_logger(self)
+
+ @abstractmethod
+ def load_dataset(self, dataset_path: str, **kwargs) -> Dataset | DatasetDict | Optional[Dataset]:
+ """
+ Load a dataset from a file.
+
+ Args:
+ dataset_path (str): The path to the dataset file.
+ split (str, optional): The split to load. Defaults to None.
+ **kwargs: Additional keyword arguments to pass to the `load_dataset` method.
+
+ Returns:
+ Union[Dataset, DatasetDict, None]: The loaded dataset.
+
+ Raises:
+ NotImplementedError: This method should be overridden by subclasses.
+ """
+ raise NotImplementedError("Subclasses should implement this!")
+
+ def preprocess_data(self, **kwargs):
+ """Load and preprocess the dataset"""
+ try:
+ if self.use_huggingface_dataset:
+ _dataset = self.load_dataset(self.huggingface_dataset, **kwargs)
+ if self.evaluate:
+ _dataset = _dataset.train_test_split(test_size=0.2)
+ self.train_dataset = _dataset["train"]
+ self.eval_dataset = _dataset["test"]
+ else:
+ self.train_dataset = _dataset["train"]
+ self.eval_dataset = None
+ elif self.evaluate:
+ train_dataset_path = os.path.join(self.input.get(), "train")
+ eval_dataset_path = os.path.join(self.input.get(), "test")
+ self.train_dataset = self.load_dataset(train_dataset_path, **kwargs)
+ self.eval_dataset = self.load_dataset(eval_dataset_path, **kwargs)
+ else:
+ self.train_dataset = self.load_dataset(self.input.get(), **kwargs)
+ self.eval_dataset = None
+ except Exception as e:
+ self.log.exception(f"Failed to preprocess data: {e}")
+ raise e
+
+ def load_models(
+ self,
+ model_name: str,
+ tokenizer_name: str,
+ model_class: str = "AutoModel",
+ tokenizer_class: str = "AutoTokenizer",
+ device_map: str | dict = "auto",
+ precision: str = "bfloat16",
+ quantization: Optional[int] = None,
+ lora_config: Optional[dict] = None,
+ use_accelerate: bool = False,
+ accelerate_no_split_module_classes: List[str] = [],
+ **kwargs,
+ ):
+ """
+ Load the model and tokenizer.
+
+ Args:
+ model_name (str): The name of the model to be loaded.
+ tokenizer_name (str, optional): The name of the tokenizer to be loaded. Defaults to None.
+ model_class (str, optional): The class of the model. Defaults to "AutoModel".
+ tokenizer_class (str, optional): The class of the tokenizer. Defaults to "AutoTokenizer".
+ device (Union[str, torch.device], optional): The device to be used. Defaults to "cuda".
+ precision (str, optional): The precision to be used. Choose from 'float32', 'float16', 'bfloat16'. Defaults to "float32".
+ quantization (Optional[int], optional): The quantization to be used. Defaults to None.
+ lora_config (Optional[dict], optional): The LoRA configuration to be used. Defaults to None.
+ use_accelerate (bool, optional): Whether to use accelerate. Defaults to False.
+ accelerate_no_split_module_classes (List[str], optional): The list of no split module classes to be used. Defaults to [].
+ **kwargs: Additional keyword arguments.
+
+ Raises:
+ ValueError: If an unsupported precision is chosen.
+
+ Returns:
+ None
+ """
+ try:
+ # Determine the torch dtype based on precision
+ if precision == "float32":
+ torch_dtype = torch.float32
+ elif precision == "float":
+ torch_dtype = torch.float
+ elif precision == "float64":
+ torch_dtype = torch.float64
+ elif precision == "double":
+ torch_dtype = torch.double
+ elif precision == "float16":
+ torch_dtype = torch.float16
+ elif precision == "bfloat16":
+ torch_dtype = torch.bfloat16
+ elif precision == "half":
+ torch_dtype = torch.half
+ elif precision == "uint8":
+ torch_dtype = torch.uint8
+ elif precision == "int8":
+ torch_dtype = torch.int8
+ elif precision == "int16":
+ torch_dtype = torch.int16
+ elif precision == "short":
+ torch_dtype = torch.short
+ elif precision == "int32":
+ torch_dtype = torch.int32
+ elif precision == "int":
+ torch_dtype = torch.int
+ elif precision == "int64":
+ torch_dtype = torch.int64
+ elif precision == "quint8":
+ torch_dtype = torch.quint8
+ elif precision == "qint8":
+ torch_dtype = torch.qint8
+ elif precision == "qint32":
+ torch_dtype = torch.qint32
+ else:
+ torch_dtype = None
+
+ peft_target_modules = []
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ else:
+ model_revision = None
+ self.model_name = model_name
+ self.log.info(f"Loading model {model_name} and branch {model_revision}")
+
+ with init_empty_weights():
+ model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ model_name, revision=model_revision, device_map=device_map
+ )
+ known_targets = [
+ v
+ for k, v in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING.items()
+ if k.lower() in model_name.lower()
+ ]
+ if len(known_targets) > 0:
+ peft_target_modules = known_targets[0]
+ else:
+ # very generic strategy, may lead to VRAM usage explosion on the wrong model erasing all advantage
+ for name, module in model.named_modules():
+ if isinstance(module, (torch.nn.Linear, torch.nn.Conv1d)) and "head" not in name:
+ name = name.split(".")[-1]
+ if name not in peft_target_modules:
+ peft_target_modules.append(name)
+ self.log.info(f"Targeting these modules for PEFT: {peft_target_modules}")
+
+ if use_accelerate:
+ if precision == "float16":
+ device_map = infer_auto_device_map(
+ model,
+ dtype="float16",
+ no_split_module_classes=accelerate_no_split_module_classes,
+ **kwargs,
+ )
+ elif precision == "bfloat16":
+ device_map = infer_auto_device_map(
+ model,
+ dtype="bfloat16",
+ no_split_module_classes=accelerate_no_split_module_classes,
+ **kwargs,
+ )
+ else:
+ device_map = infer_auto_device_map(
+ model,
+ no_split_module_classes=accelerate_no_split_module_classes,
+ **kwargs,
+ )
+ self.log.info(f"Inferred device map {device_map}")
+
+ # Create the LoRA config for PEFT
+ if lora_config:
+ if len(peft_target_modules) > 0:
+ lora_config = LoraConfig(target_modules=peft_target_modules, **lora_config)
+ else:
+ lora_config = LoraConfig(**lora_config)
+ self.peft_config = lora_config
+ # you cannot fine-tune quantized models without LoRA
+ if quantization and not lora_config:
+ lora_config = {
+ "r": 16,
+ "lora_alpha": 32,
+ "lora_dropout": 0.05,
+ "bias": "none",
+ "task_type": "CAUSAL_LM",
+ }
+ self.peft_config = LoraConfig(target_modules=peft_target_modules, **lora_config)
+ if lora_config:
+ self.log.info(f"LoRA config: {self.peft_config}")
+
+ # Load model and tokenizer
+ if quantization == 8:
+ # Use AutoConfig to automatically load the configuration
+ if self.model_name.lower() == "local": # type: ignore
+ self.log.info(f"Loading local model {model_class} : {self.input.get()}")
+ if not self.config: # type: ignore
+ self.config = AutoConfig.from_pretrained(os.path.join(self.input.get(), "/model"))
+ self.model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ os.path.join(self.input.get(), "/model"),
+ device_map=device_map,
+ torch_dtype=torch_dtype,
+ load_in_8bit=True,
+ config=self.config,
+ **kwargs,
+ )
+ else:
+ self.log.info(f"Loading from huggingface hub: {model_class} : {model_name}")
+ if not self.config:
+ self.config = AutoConfig.from_pretrained(self.model_name)
+ self.model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ self.model_name,
+ revision=model_revision,
+ device_map=device_map,
+ torch_dtype=torch_dtype,
+ load_in_8bit=True,
+ config=self.config,
+ **kwargs,
+ )
+ elif quantization == 4:
+ # Use AutoConfig to automatically load the configuration
+ if self.model_name.lower() == "local": # type: ignore
+ self.log.info(f"Loading local model {model_class} : {self.input.get()}")
+ if not self.config:
+ self.config = AutoConfig.from_pretrained(os.path.join(self.input.get(), "/model"))
+ self.model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ os.path.join(self.input.get(), "/model"),
+ device_map=device_map,
+ torch_dtype=torch_dtype,
+ load_in_4bit=True,
+ config=self.config,
+ **kwargs,
+ )
+ else:
+ self.log.info(f"Loading from huggingface hub: {model_class} : {model_name}")
+ if not self.config:
+ self.config = AutoConfig.from_pretrained(self.model_name)
+ self.model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ self.model_name,
+ revision=model_revision,
+ device_map=device_map,
+ torch_dtype=torch_dtype,
+ load_in_4bit=True,
+ config=self.config,
+ **kwargs,
+ )
+ else:
+ # Use AutoConfig to automatically load the configuration
+ if self.model_name.lower() == "local": # type: ignore
+ self.log.info(f"Loading local model {model_class} : {self.input.get()}")
+ if not self.config:
+ self.config = AutoConfig.from_pretrained(os.path.join(self.input.get(), "/model"))
+ self.model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ os.path.join(self.input.get(), "/model"),
+ device_map=device_map,
+ torch_dtype=torch_dtype,
+ config=self.config,
+ **kwargs,
+ )
+ else:
+ self.log.info(f"Loading from huggingface hub: {model_class} : {model_name}")
+ if not self.config:
+ self.config = AutoConfig.from_pretrained(self.model_name)
+ self.model = getattr(__import__("transformers"), str(model_class)).from_pretrained(
+ model_name,
+ revision=model_revision,
+ device_map=device_map,
+ torch_dtype=torch_dtype,
+ config=self.config,
+ **kwargs,
+ )
+
+ if ":" in tokenizer_name:
+ tokenizer_revision = tokenizer_name.split(":")[1]
+ tokenizer_name = tokenizer_name.split(":")[0]
+ else:
+ tokenizer_revision = None
+ self.tokenizer_name = tokenizer_name
+ self.tokenizer_revision = tokenizer_revision
+
+ if tokenizer_name.lower() == "local": # type: ignore
+ self.log.info(f"Loading local tokenizer : {tokenizer_class} : {self.input.get()}")
+ self.tokenizer = getattr(__import__("transformers"), str(tokenizer_class)).from_pretrained(
+ os.path.join(self.input.get(), "/model")
+ )
+ else:
+ self.log.info(f"Loading tokenizer from huggingface hub: {tokenizer_class} : {tokenizer_name}")
+ self.tokenizer = getattr(__import__("transformers"), str(tokenizer_class)).from_pretrained(
+ tokenizer_name, revision=tokenizer_revision
+ )
+ except Exception as e:
+ self.log.exception(f"Failed to load model: {e}")
+ raise
+
+ def upload_to_hf_hub(
+ self,
+ hf_repo_id: Optional[str] = None,
+ hf_commit_message: Optional[str] = None,
+ hf_token: Optional[str] = None,
+ hf_private: Optional[str] = None,
+ hf_create_pr: Optional[str] = None,
+ ):
+ """Upload the model and tokenizer to Hugging Face Hub."""
+ try:
+ if self.model:
+ self.model.to("cpu").push_to_hub(
+ repo_id=hf_repo_id if hf_repo_id else self.hf_repo_id, # type: ignore
+ commit_message=hf_commit_message if hf_commit_message else self.hf_commit_message, # type: ignore
+ token=hf_token if hf_token else self.hf_token, # type: ignore
+ private=hf_private if hf_private else self.hf_private, # type: ignore
+ create_pr=hf_create_pr if hf_create_pr else self.hf_create_pr, # type: ignore
+ )
+ if self.tokenizer:
+ self.tokenizer.push_to_hub(
+ repo_id=hf_repo_id if hf_repo_id else self.hf_repo_id, # type: ignore
+ commit_message=hf_commit_message if hf_commit_message else self.hf_commit_message, # type: ignore
+ token=hf_token if hf_token else self.hf_token, # type: ignore
+ private=hf_private if hf_private else self.hf_private, # type: ignore
+ create_pr=hf_create_pr if hf_create_pr else self.hf_create_pr, # type: ignore
+ )
+ except Exception as e:
+ self.log.exception(f"Failed to upload model to huggingface hub: {e}")
+ raise
+
+ def compute_metrics(self, eval_pred: EvalPrediction) -> Optional[Dict[str, float]] | Dict[str, float]:
+ """
+ Compute metrics for evaluation. This class implements a simple classification evaluation, tasks should ideally override this.
+
+ Args:
+ eval_pred (EvalPrediction): The evaluation predictions.
+
+ Returns:
+ dict: The computed metrics.
+ """
+ predictions, labels = eval_pred
+ predictions = predictions[0] if isinstance(predictions, tuple) else predictions
+ labels = labels[0] if isinstance(labels, tuple) else labels
+ predictions = np.argmax(predictions, axis=1)
+
+ return {
+ "accuracy": accuracy_score(labels, predictions),
+ "precision": precision_recall_fscore_support(labels, predictions, average="binary")[0],
+ "recall": precision_recall_fscore_support(labels, predictions, average="binary")[1],
+ "f1": precision_recall_fscore_support(labels, predictions, average="binary")[2],
+ }
+
+ def fine_tune(
+ self,
+ model_name: str,
+ tokenizer_name: str,
+ num_train_epochs: int,
+ per_device_batch_size: int,
+ model_class: str = "AutoModel",
+ tokenizer_class: str = "AutoTokenizer",
+ device_map: str | dict = "auto",
+ precision: str = "bfloat16",
+ quantization: Optional[int] = None,
+ lora_config: Optional[dict] = None,
+ use_accelerate: bool = False,
+ use_trl: bool = False,
+ accelerate_no_split_module_classes: List[str] = [],
+ compile: bool = False,
+ evaluate: bool = False,
+ save_steps: int = 500,
+ save_total_limit: Optional[int] = None,
+ load_best_model_at_end: bool = False,
+ metric_for_best_model: Optional[str] = None,
+ greater_is_better: Optional[bool] = None,
+ map_data: Optional[Callable] = None,
+ use_huggingface_dataset: bool = False,
+ huggingface_dataset: str = "",
+ hf_repo_id: Optional[str] = None,
+ hf_commit_message: Optional[str] = None,
+ hf_token: Optional[str] = None,
+ hf_private: bool = True,
+ hf_create_pr: bool = False,
+ notification_email: str = "",
+ learning_rate: float = 1e-5,
+ **kwargs,
+ ):
+ """
+ Fine-tunes a pre-trained Hugging Face model.
+
+ Args:
+ model_name (str): The name of the pre-trained model.
+ tokenizer_name (str): The name of the pre-trained tokenizer.
+ num_train_epochs (int): The total number of training epochs to perform.
+ per_device_batch_size (int): The batch size per device during training.
+ model_class (str, optional): The model class to use. Defaults to "AutoModel".
+ tokenizer_class (str, optional): The tokenizer class to use. Defaults to "AutoTokenizer".
+ device_map (str | dict, optional): The device map for distributed training. Defaults to "auto".
+ precision (str, optional): The precision to use for training. Defaults to "bfloat16".
+ quantization (int, optional): The quantization level to use for training. Defaults to None.
+ lora_config (dict, optional): Configuration for PEFT LoRA optimization. Defaults to None.
+ use_accelerate (bool, optional): Whether to use accelerate for distributed training. Defaults to False.
+ use_trl (bool, optional): Whether to use TRL for training. Defaults to False.
+ accelerate_no_split_module_classes (List[str], optional): The module classes to not split during distributed training. Defaults to [].
+ evaluate (bool, optional): Whether to evaluate the model after training. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ save_steps (int, optional): Number of steps between checkpoints. Defaults to 500.
+ save_total_limit (Optional[int], optional): Maximum number of checkpoints to keep. Older checkpoints are deleted. Defaults to None.
+ load_best_model_at_end (bool, optional): Whether to load the best model (according to evaluation) at the end of training. Defaults to False.
+ metric_for_best_model (Optional[str], optional): The metric to use to compare models. Defaults to None.
+ greater_is_better (Optional[bool], optional): Whether a larger value of the metric indicates a better model. Defaults to None.
+ use_huggingface_dataset (bool, optional): Whether to load a dataset from huggingface hub.
+ huggingface_dataset (str, optional): The huggingface dataset to use.
+ map_data (Callable, optional): A function to map data before training. Defaults to None.
+ hf_repo_id (str, optional): The Hugging Face repo ID. Defaults to None.
+ hf_commit_message (str, optional): The Hugging Face commit message. Defaults to None.
+ hf_token (str, optional): The Hugging Face token. Defaults to None.
+ hf_private (bool, optional): Whether to make the repo private. Defaults to True.
+ hf_create_pr (bool, optional): Whether to create a pull request. Defaults to False.
+ notification_email (str, optional): Whether to notify after job is complete. Defaults to None.
+ learning_rate (float, optional): Learning rate for backpropagation.
+ **kwargs: Additional keyword arguments to pass to the model.
+
+ Returns:
+ None
+ """
+ try:
+ # Save everything
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.num_train_epochs = num_train_epochs
+ self.per_device_batch_size = per_device_batch_size
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.device_map = device_map
+ self.precision = precision
+ self.quantization = quantization
+ self.lora_config = lora_config # type: ignore
+ self.use_accelerate = use_accelerate
+ self.use_trl = use_trl
+ self.accelerate_no_split_module_classes = accelerate_no_split_module_classes
+ self.evaluate = evaluate
+ self.use_huggingface_dataset = use_huggingface_dataset
+ self.huggingface_dataset = huggingface_dataset
+ self.hf_repo_id = hf_repo_id
+ self.hf_commit_message = hf_commit_message
+ self.hf_token = hf_token
+ self.hf_private = hf_private
+ self.hf_create_pr = hf_create_pr
+ self.map_data = map_data
+ self.notification_email = notification_email
+ self.learning_rate = learning_rate
+ self.config = None
+
+ model_kwargs = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_kwargs = model_kwargs
+
+ self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ device_map=self.device_map,
+ precision=self.precision,
+ quantization=self.quantization,
+ lora_config=self.lora_config,
+ use_accelerate=self.use_accelerate,
+ accelerate_no_split_module_classes=self.accelerate_no_split_module_classes,
+ **model_kwargs,
+ )
+
+ if self.tokenizer and not self.tokenizer.pad_token:
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+ self.model.config.pad_token_id = self.tokenizer.eos_token_id
+
+ # Load dataset
+ dataset_kwargs = {k.replace("data_", ""): v for k, v in kwargs.items() if "data_" in k}
+ self.dataset_kwargs = dataset_kwargs
+ self.preprocess_data(**dataset_kwargs)
+
+ # Separate training and evaluation arguments
+ trainer_kwargs = {k.replace("trainer_", ""): v for k, v in kwargs.items() if "trainer_" in k}
+ self.trainer_kwargs = trainer_kwargs
+ training_kwargs = {k.replace("training_", ""): v for k, v in kwargs.items() if "training_" in k}
+ self.training_kwargs = training_kwargs
+
+ # Create training arguments
+ training_args = TrainingArguments(
+ output_dir=os.path.join(self.output.output_folder, "model"),
+ num_train_epochs=num_train_epochs,
+ per_device_train_batch_size=per_device_batch_size,
+ per_device_eval_batch_size=per_device_batch_size,
+ save_steps=save_steps,
+ save_total_limit=save_total_limit,
+ load_best_model_at_end=load_best_model_at_end,
+ metric_for_best_model=metric_for_best_model,
+ greater_is_better=greater_is_better,
+ dataloader_num_workers=4,
+ learning_rate=self.learning_rate,
+ **training_kwargs,
+ )
+
+ # Add adapters to the model for fine-tuning
+ if self.lora_config and not use_trl:
+ self.model.enable_input_require_grads()
+ self.model = get_peft_model(self.model, peft_config=self.peft_config)
+
+ if compile:
+ self.model = torch.compile(self.model)
+
+ # Create trainer
+ if use_trl:
+ self.model = get_peft_model(self.model, peft_config=self.peft_config)
+ trainer = SFTTrainer(
+ model=self.model,
+ args=training_args,
+ train_dataset=self.train_dataset,
+ eval_dataset=self.eval_dataset,
+ tokenizer=self.tokenizer,
+ compute_metrics=self.compute_metrics,
+ data_collator=self.data_collator if hasattr(self, "data_collator") else None,
+ peft_config=self.peft_config,
+ **trainer_kwargs,
+ )
+ else:
+ trainer = Trainer(
+ model=self.model,
+ args=training_args,
+ train_dataset=self.train_dataset,
+ eval_dataset=self.eval_dataset,
+ tokenizer=self.tokenizer,
+ compute_metrics=self.compute_metrics,
+ data_collator=self.data_collator if hasattr(self, "data_collator") else None,
+ **trainer_kwargs,
+ )
+
+ # Train the model
+ trainer.train()
+ trainer.save_model(self.output.output_folder)
+
+ if self.evaluate:
+ eval_result = trainer.evaluate()
+ self.log.info(f"Evaluation results: {eval_result}")
+
+ # Save the model configuration to Hugging Face Hub if hf_repo_id is not None
+ if self.hf_repo_id and self.config:
+ self.config.save_pretrained(os.path.join(self.output.output_folder, "model"))
+ self.upload_to_hf_hub()
+ except Exception as e:
+ self.log.exception(f"Failed to fine tune model: {e}")
+ self.state.set_state(self.id, {"success": False, "exception": str(e)})
+ raise
+ self.state.set_state(self.id, {"success": True})
+
+ self.done()
+
+ def done(self):
+ if self.notification_email:
+ self.output.flush()
+ send_fine_tuning_email(
+ recipient=self.notification_email, bucket_name=self.output.bucket, prefix=self.output.s3_folder
+ )
diff --git a/geniusrise_text/base/tests/test_api.py b/geniusrise_text/base/tests/test_api.py
new file mode 100644
index 0000000..77d1dd3
--- /dev/null
+++ b/geniusrise_text/base/tests/test_api.py
@@ -0,0 +1,211 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import os
+
+import pytest
+import torch
+
+# import transformers
+from geniusrise.core import BatchInput, BatchOutput, InMemoryState
+
+from geniusrise_text.base.api import TextAPI
+
+
+@pytest.fixture(
+ params=[
+ # model_name, model_class, tokenizer_class, use_cuda, precision, quantization, device_map, max_memory, torchscript
+ # fmt: off
+ ("gpt2", "AutoModelForCausalLM", "AutoTokenizer", True, "float16", 0, None, None, False),
+ ("gpt2", "AutoModelForCausalLM", "AutoTokenizer", False, "float32", 0, None, None, False),
+ ("bigscience/bloom-560m", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 0, None, None, False),
+ ("meta-llama/Llama-2-7b-hf", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 4, None, None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 0, "cuda:0", None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 4, "cuda:0", None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 8, "cuda:0", None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 0, "auto", None, True),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 4, "auto", None, True),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 8, "auto", None, True),
+ ("TheBloke/Mistral-7B-v0.1-GPTQ:gptq-4bit-32g-actorder_True", "AutoModelForCausalLM", "AutoTokenizer", True, "float16", None, "cuda:0", None, False),
+ # fmt: on
+ ]
+)
+def model_config(request):
+ return request.param
+
+
+# Fixtures to initialize TextAPI instance
+@pytest.fixture
+def hfa():
+ input_dir = "./input_dir"
+ output_dir = "./output_dir"
+
+ input = BatchInput(input_dir, "geniusrise-test", "api_input")
+ output = BatchOutput(output_dir, "geniusrise-test", "api_output")
+ state = InMemoryState()
+
+ hfa = TextAPI(
+ input=input,
+ output=output,
+ state=state,
+ )
+ yield hfa # provide the fixture value
+
+ # cleanup
+ if os.path.exists(input_dir):
+ os.rmdir(input_dir)
+ if os.path.exists(output_dir):
+ os.rmdir(output_dir)
+
+
+def test_load_models(hfa, model_config):
+ (
+ model_name,
+ model_class,
+ tokenizer_class,
+ use_cuda,
+ precision,
+ quantization,
+ device_map,
+ max_memory,
+ torchscript,
+ ) = model_config
+
+ if ":" in model_name:
+ _model_name = model_name
+ model_revision = _model_name.split(":")[1]
+ model_name = _model_name.split(":")[0]
+ tokenizer_revision = _model_name.split(":")[1]
+ tokenizer_name = _model_name.split(":")[0]
+ else:
+ model_revision = None
+ tokenizer_revision = None
+
+ model, tokenizer = hfa.load_models(
+ model_name=model_name,
+ model_revision=model_revision,
+ tokenizer_name=model_name,
+ tokenizer_revision=tokenizer_revision,
+ model_class=model_class,
+ tokenizer_class=tokenizer_class,
+ use_cuda=use_cuda,
+ precision=precision,
+ quantization=quantization,
+ device_map=device_map,
+ max_memory=max_memory,
+ torchscript=torchscript,
+ )
+ assert model is not None
+ assert tokenizer is not None
+ assert len(list(model.named_modules())) > 0
+
+ del model
+ del tokenizer
+ torch.cuda.empty_cache()
+
+
+# Define strategies and associated parameters
+strategies = {
+ "generate": {},
+ "greedy_search": {},
+ "beam_search": {"num_beams": 4},
+ "beam_sample": {"num_beams": 4, "temperature": 0.7, "top_k": 20},
+ "group_beam_search": {"num_beams": 4, "num_beam_groups": 2},
+}
+
+# Define other parameters
+length_params = {
+ "max_length": [20, 30],
+ "min_length": [0, 10],
+ "early_stopping": [False, True],
+}
+gen_strategy_params = {
+ "do_sample": [False, True],
+}
+logit_params = {
+ "temperature": [1.0, 0.7],
+ "top_k": [50, 20],
+ "top_p": [1.0, 0.9],
+ "repetition_penalty": [1.0, 1.5],
+ "length_penalty": [1.0, 0.5],
+ "no_repeat_ngram_size": [0, 2],
+}
+# Merge all the parameters into one dictionary for itertools.product
+all_params = {**length_params, **gen_strategy_params, **logit_params}
+
+
+@pytest.mark.parametrize("strategy", list(strategies.keys()))
+def test_generate_strategies(hfa, model_config, strategy):
+ (
+ model_name,
+ model_class,
+ tokenizer_class,
+ use_cuda,
+ precision,
+ quantization,
+ device_map,
+ max_memory,
+ torchscript,
+ ) = model_config
+
+ if ":" in model_name:
+ _model_name = model_name
+ model_revision = _model_name.split(":")[1]
+ model_name = _model_name.split(":")[0]
+ tokenizer_revision = _model_name.split(":")[1]
+ tokenizer_name = _model_name.split(":")[0]
+ else:
+ model_revision = None
+ tokenizer_revision = None
+
+ model, tokenizer = hfa.load_models(
+ model_name=model_name,
+ model_revision=model_revision,
+ tokenizer_name=model_name,
+ tokenizer_revision=tokenizer_revision,
+ model_class=model_class,
+ tokenizer_class=tokenizer_class,
+ use_cuda=use_cuda,
+ precision=precision,
+ quantization=quantization,
+ device_map=device_map,
+ max_memory=max_memory,
+ torchscript=torchscript,
+ )
+ hfa.model = model
+ hfa.tokenizer = tokenizer
+
+ # Strategy-specific params
+ strategy_params = strategies[strategy]
+
+ # All possible combinations for the current strategy
+ param_combinations = [
+ {**dict(zip(all_params.keys(), values)), **strategy_params}
+ for values in itertools.product(*all_params.values())
+ ]
+
+ for param_set in param_combinations:
+ generated_text = hfa.generate(
+ prompt="Once upon a time", decoding_strategy=strategy, **param_set # Unpack params into function arguments
+ )
+ assert generated_text is not None
+ assert isinstance(generated_text, str)
+ break
+
+ # Cleanup
+ del model
+ del tokenizer
+ torch.cuda.empty_cache()
diff --git a/geniusrise_text/base/tests/test_bulk.py b/geniusrise_text/base/tests/test_bulk.py
new file mode 100644
index 0000000..039ca3d
--- /dev/null
+++ b/geniusrise_text/base/tests/test_bulk.py
@@ -0,0 +1,211 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import os
+
+import pytest
+import torch
+
+# import transformers
+from geniusrise.core import BatchInput, BatchOutput, InMemoryState
+
+from geniusrise_text.base.bulk import TextBulk
+
+
+@pytest.fixture(
+ params=[
+ # model_name, model_class, tokenizer_class, use_cuda, precision, quantization, device_map, max_memory, torchscript
+ # fmt: off
+ ("gpt2", "AutoModelForCausalLM", "AutoTokenizer", True, "float16", 0, None, None, False),
+ ("gpt2", "AutoModelForCausalLM", "AutoTokenizer", False, "float32", 0, None, None, False),
+ ("bigscience/bloom-560m", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 0, None, None, False),
+ ("meta-llama/Llama-2-7b-hf", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 4, None, None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 0, "cuda:0", None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 4, "cuda:0", None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 8, "cuda:0", None, False),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 0, "auto", None, True),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 4, "auto", None, True),
+ ("mistralai/Mistral-7B-v0.1", "AutoModelForCausalLM", "AutoTokenizer", True, "bfloat16", 8, "auto", None, True),
+ ("TheBloke/Mistral-7B-v0.1-GPTQ:gptq-4bit-32g-actorder_True", "AutoModelForCausalLM", "AutoTokenizer", True, "float16", None, "cuda:0", None, False),
+ # fmt: on
+ ]
+)
+def model_config(request):
+ return request.param
+
+
+# Fixtures to initialize TextBulk instance
+@pytest.fixture
+def hfa():
+ input_dir = "./input_dir"
+ output_dir = "./output_dir"
+
+ input = BatchInput(input_dir, "geniusrise-test", "api_input")
+ output = BatchOutput(output_dir, "geniusrise-test", "api_output")
+ state = InMemoryState()
+
+ hfa = TextBulk(
+ input=input,
+ output=output,
+ state=state,
+ )
+ yield hfa # provide the fixture value
+
+ # cleanup
+ if os.path.exists(input_dir):
+ os.rmdir(input_dir)
+ if os.path.exists(output_dir):
+ os.rmdir(output_dir)
+
+
+def test_load_models(hfa, model_config):
+ (
+ model_name,
+ model_class,
+ tokenizer_class,
+ use_cuda,
+ precision,
+ quantization,
+ device_map,
+ max_memory,
+ torchscript,
+ ) = model_config
+
+ if ":" in model_name:
+ _model_name = model_name
+ model_revision = _model_name.split(":")[1]
+ model_name = _model_name.split(":")[0]
+ tokenizer_revision = _model_name.split(":")[1]
+ tokenizer_name = _model_name.split(":")[0]
+ else:
+ model_revision = None
+ tokenizer_revision = None
+
+ model, tokenizer = hfa.load_models(
+ model_name=model_name,
+ model_revision=model_revision,
+ tokenizer_name=model_name,
+ tokenizer_revision=tokenizer_revision,
+ model_class=model_class,
+ tokenizer_class=tokenizer_class,
+ use_cuda=use_cuda,
+ precision=precision,
+ quantization=quantization,
+ device_map=device_map,
+ max_memory=max_memory,
+ torchscript=torchscript,
+ )
+ assert model is not None
+ assert tokenizer is not None
+ assert len(list(model.named_modules())) > 0
+
+ del model
+ del tokenizer
+ torch.cuda.empty_cache()
+
+
+# Define strategies and associated parameters
+strategies = {
+ "generate": {},
+ "greedy_search": {},
+ "beam_search": {"num_beams": 4},
+ "beam_sample": {"num_beams": 4, "temperature": 0.7, "top_k": 20},
+ "group_beam_search": {"num_beams": 4, "num_beam_groups": 2},
+}
+
+# Define other parameters
+length_params = {
+ "max_length": [20, 30],
+ "min_length": [0, 10],
+ "early_stopping": [False, True],
+}
+gen_strategy_params = {
+ "do_sample": [False, True],
+}
+logit_params = {
+ "temperature": [1.0, 0.7],
+ "top_k": [50, 20],
+ "top_p": [1.0, 0.9],
+ "repetition_penalty": [1.0, 1.5],
+ "length_penalty": [1.0, 0.5],
+ "no_repeat_ngram_size": [0, 2],
+}
+# Merge all the parameters into one dictionary for itertools.product
+all_params = {**length_params, **gen_strategy_params, **logit_params}
+
+
+@pytest.mark.parametrize("strategy", list(strategies.keys()))
+def test_generate_strategies(hfa, model_config, strategy):
+ (
+ model_name,
+ model_class,
+ tokenizer_class,
+ use_cuda,
+ precision,
+ quantization,
+ device_map,
+ max_memory,
+ torchscript,
+ ) = model_config
+
+ if ":" in model_name:
+ _model_name = model_name
+ model_revision = _model_name.split(":")[1]
+ model_name = _model_name.split(":")[0]
+ tokenizer_revision = _model_name.split(":")[1]
+ tokenizer_name = _model_name.split(":")[0]
+ else:
+ model_revision = None
+ tokenizer_revision = None
+
+ model, tokenizer = hfa.load_models(
+ model_name=model_name,
+ model_revision=model_revision,
+ tokenizer_name=model_name,
+ tokenizer_revision=tokenizer_revision,
+ model_class=model_class,
+ tokenizer_class=tokenizer_class,
+ use_cuda=use_cuda,
+ precision=precision,
+ quantization=quantization,
+ device_map=device_map,
+ max_memory=max_memory,
+ torchscript=torchscript,
+ )
+ hfa.model = model
+ hfa.tokenizer = tokenizer
+
+ # Strategy-specific params
+ strategy_params = strategies[strategy]
+
+ # All possible combinations for the current strategy
+ param_combinations = [
+ {**dict(zip(all_params.keys(), values)), **strategy_params}
+ for values in itertools.product(*all_params.values())
+ ]
+
+ for param_set in param_combinations:
+ generated_text = hfa.generate(
+ prompt="Once upon a time", decoding_strategy=strategy, **param_set # Unpack params into function arguments
+ )
+ assert generated_text is not None
+ assert isinstance(generated_text, str)
+ break
+
+ # Cleanup
+ del model
+ del tokenizer
+ torch.cuda.empty_cache()
diff --git a/geniusrise_text/base/tests/test_fine_tune.py b/geniusrise_text/base/tests/test_fine_tune.py
new file mode 100644
index 0000000..4ea1175
--- /dev/null
+++ b/geniusrise_text/base/tests/test_fine_tune.py
@@ -0,0 +1,257 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from datasets import load_dataset
+from geniusrise.core import BatchInput, BatchOutput, InMemoryState
+from transformers import DataCollatorForLanguageModeling, EvalPrediction
+
+from geniusrise_text.base import TextFineTuner
+
+# SEQ_CLS = "SEQ_CLS"
+# SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM"
+# CAUSAL_LM = "CAUSAL_LM"
+# TOKEN_CLS = "TOKEN_CLS"
+# QUESTION_ANS = "QUESTION_ANS"
+# FEATURE_EXTRACTION = "FEATURE_EXTRACTION"
+
+lora_config = {
+ "r": 16,
+ "lora_alpha": 32,
+ "lora_dropout": 0.05,
+ "bias": "none",
+ "task_type": "CAUSAL_LM",
+}
+
+
+class TestTextFineTuner(TextFineTuner):
+ def load_dataset(self, dataset_path, **kwargs):
+ dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1%]") # Adjust the split as needed
+ dataset = dataset.map(
+ lambda examples: self.tokenizer(
+ examples["text"],
+ truncation=True,
+ padding="max_length",
+ max_length=512,
+ ),
+ batched=True,
+ )
+ return dataset
+
+ def data_collator(self, examples):
+ return DataCollatorForLanguageModeling(self.tokenizer, mlm=False)(examples)
+
+
+@pytest.fixture
+def bolt():
+ input_dir = tempfile.mkdtemp()
+ output_dir = tempfile.mkdtemp()
+
+ input = BatchInput(input_dir, "geniusrise-test", "test-🤗-input")
+ output = BatchOutput(output_dir, "geniusrise-test", "test-🤗-output")
+ state = InMemoryState()
+
+ return TestTextFineTuner(
+ input=input,
+ output=output,
+ state=state,
+ evaluate=False,
+ )
+
+
+def test_bolt_init(bolt):
+ assert bolt.input is not None
+ assert bolt.output is not None
+ assert bolt.state is not None
+
+
+def test_load_dataset(bolt):
+ bolt.model_name = "bert-base-uncased"
+ bolt.tokenizer_name = "bert-base-uncased"
+ bolt.model_class = "AutoModelForCausalLM"
+ bolt.tokenizer_class = "BertTokenizer"
+ bolt.load_models(
+ model_name=bolt.model_name,
+ tokenizer_name=bolt.tokenizer_name,
+ model_class=bolt.model_class,
+ tokenizer_class=bolt.tokenizer_class,
+ device_map=None,
+ )
+ dataset = bolt.load_dataset("fake_path")
+ assert dataset is not None
+ assert len(dataset) >= 100
+
+ del bolt.model
+ del bolt.tokenizer
+ torch.cuda.empty_cache()
+
+
+def test_fine_tune(bolt):
+ bolt.fine_tune(
+ model_name="bert-base-uncased",
+ tokenizer_name="bert-base-uncased",
+ num_train_epochs=1,
+ per_device_batch_size=2,
+ model_class="AutoModelForCausalLM",
+ tokenizer_class="BertTokenizer",
+ evaluate=False,
+ device_map=None,
+ )
+ bolt.upload_to_hf_hub(
+ hf_repo_id="ixaxaar/geniusrise-hf-base-test-repo",
+ hf_commit_message="testing base fine tuner",
+ hf_token=os.getenv("HUGGINGFACE_ACCESS_TOKEN"),
+ hf_private=False,
+ hf_create_pr=True,
+ )
+
+ # Check that model files are created in the output directory
+ assert os.path.isfile(os.path.join(bolt.output.output_folder, "model", "pytorch_model.bin"))
+ assert os.path.isfile(os.path.join(bolt.output.output_folder, "model", "config.json"))
+ assert os.path.isfile(os.path.join(bolt.output.output_folder, "model", "training_args.bin"))
+
+ del bolt.model
+ del bolt.tokenizer
+ torch.cuda.empty_cache()
+
+
+def test_compute_metrics(bolt):
+ # Mocking an EvalPrediction object
+ logits = np.array([[0.6, 0.4], [0.4, 0.6]])
+ labels = np.array([0, 1])
+ eval_pred = EvalPrediction(predictions=logits, label_ids=labels)
+
+ metrics = bolt.compute_metrics(eval_pred)
+
+ assert "accuracy" in metrics
+ assert "precision" in metrics
+ assert "recall" in metrics
+ assert "f1" in metrics
+
+
+models = {
+ "small": "bigscience/bloom-560m",
+ "medium": "meta-llama/Llama-2-7b-hf",
+ "large": "mistralai/Mistral-7B-v0.1",
+ "4-bit": "TheBloke/Mistral-7B-v0.1-GPTQ:gptq-4bit-32g-actorder_True",
+ "8-bit": "TheBloke/OpenHermes-2-Mistral-7B-GPTQ:gptq-8bit-128g-actorder_True",
+}
+
+
+@pytest.mark.parametrize(
+ "model, precision, quantization, lora_config, use_accelerate",
+ [
+ # small
+ (models["small"], "float16", None, None, False),
+ (models["small"], "float16", None, None, True),
+ (models["small"], "float16", None, lora_config, False),
+ (models["small"], "float16", None, lora_config, True),
+ (models["small"], "float32", None, None, False),
+ (models["small"], "float32", None, None, True),
+ (models["small"], "float32", None, lora_config, False),
+ (models["small"], "float32", None, lora_config, True),
+ (models["small"], "bfloat16", None, None, False),
+ (models["small"], "bfloat16", None, None, True),
+ (models["small"], "bfloat16", None, lora_config, False),
+ (models["small"], "bfloat16", None, lora_config, True),
+ # small - 4bit
+ (models["small"], "float16", 4, lora_config, False),
+ (models["small"], "float16", 4, lora_config, True),
+ (models["small"], "float32", 4, lora_config, False),
+ (models["small"], "float32", 4, lora_config, True),
+ (models["small"], "bfloat16", 4, lora_config, False),
+ (models["small"], "bfloat16", 4, lora_config, True),
+ # small - 8 bit
+ (models["small"], "float16", 8, lora_config, False),
+ (models["small"], "float16", 8, lora_config, True),
+ (models["small"], "float32", 8, lora_config, False),
+ (models["small"], "float32", 8, lora_config, True),
+ (models["small"], "bfloat16", 8, lora_config, False),
+ (models["small"], "bfloat16", 8, lora_config, True),
+ # large
+ (models["large"], "bfloat16", 4, lora_config, False),
+ (models["large"], "bfloat16", 4, lora_config, True),
+ (models["large"], "float16", 4, lora_config, False),
+ (models["large"], "float16", 4, lora_config, True),
+ (models["large"], "float32", 4, lora_config, False),
+ (models["large"], "float32", 4, lora_config, True),
+ # 4 bit
+ (models["4-bit"], "float16", None, lora_config, False),
+ # 8 bit
+ (models["8-bit"], "float16", None, lora_config, False),
+ (models["8-bit"], "float16", None, lora_config, True),
+ ],
+)
+def test_fine_tune_options(bolt, model, precision, quantization, lora_config, use_accelerate):
+ use_trl = False
+
+ if use_trl:
+ bolt.fine_tune(
+ model_name=model,
+ tokenizer_name=model,
+ model_class="AutoModelForCausalLM",
+ tokenizer_class="AutoTokenizer",
+ num_train_epochs=1,
+ per_device_batch_size=2,
+ precision=precision,
+ quantization=quantization,
+ lora_config=lora_config,
+ use_accelerate=use_accelerate,
+ device_map="auto" if "GPTQ" in model else None,
+ trainer_packing=False if lora_config is not None else None,
+ trainer_dataset_text_field="text" if lora_config is not None else None,
+ )
+ else:
+ bolt.fine_tune(
+ model_name=model,
+ tokenizer_name=model,
+ model_class="AutoModelForCausalLM",
+ tokenizer_class="AutoTokenizer",
+ num_train_epochs=1,
+ per_device_batch_size=2,
+ precision=precision,
+ quantization=quantization,
+ lora_config=lora_config,
+ use_accelerate=use_accelerate,
+ device_map="auto" if "GPTQ" in model else None,
+ )
+
+ # Verify the model has been fine-tuned by checking the existence of model files
+ assert os.path.exists(os.path.join(bolt.output.output_folder, "model", "pytorch_model.bin")) or os.path.exists(
+ os.path.join(bolt.output.output_folder, "model", "adapter_model.bin")
+ )
+ assert os.path.exists(os.path.join(bolt.output.output_folder, "model", "config.json")) or os.path.exists(
+ os.path.join(bolt.output.output_folder, "model", "adapter_config.json")
+ )
+ assert os.path.exists(os.path.join(bolt.output.output_folder, "model", "training_args.bin"))
+
+ # Clear the output directory for the next test
+ try:
+ os.remove(os.path.join(bolt.output.output_folder, "model", "pytorch_model.bin"))
+ os.remove(os.path.join(bolt.output.output_folder, "model", "adapter_model.bin"))
+ os.remove(os.path.join(bolt.output.output_folder, "model", "config.json"))
+ os.remove(os.path.join(bolt.output.output_folder, "model", "adapter_config.json"))
+ os.remove(os.path.join(bolt.output.output_folder, "model", "training_args.bin"))
+ except Exception as _:
+ pass
+
+ del bolt.model
+ del bolt.tokenizer
+ torch.cuda.empty_cache()
diff --git a/geniusrise_text/base/util.py b/geniusrise_text/base/util.py
new file mode 100644
index 0000000..19ccd62
--- /dev/null
+++ b/geniusrise_text/base/util.py
@@ -0,0 +1,46 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
+ "t5": ["q", "v"],
+ "mt5": ["q", "v"],
+ "bart": ["q_proj", "v_proj"],
+ "gpt2": ["c_attn"],
+ "bloom": ["query_key_value"],
+ "blip-2": ["q", "v", "q_proj", "v_proj"],
+ "opt": ["q_proj", "v_proj"],
+ "gptj": ["q_proj", "v_proj"],
+ "gpt_neox": ["query_key_value"],
+ "gpt_neo": ["q_proj", "v_proj"],
+ "bert": ["query", "value"],
+ "roberta": ["query", "value"],
+ "xlm-roberta": ["query", "value"],
+ "electra": ["query", "value"],
+ "deberta-v2": ["query_proj", "value_proj"],
+ "deberta": ["in_proj"],
+ "layoutlm": ["query", "value"],
+ "llama": ["q_proj", "v_proj"],
+ "chatglm": ["query_key_value"],
+ "gpt_bigcode": ["c_attn"],
+ "mpt": ["Wqkv"],
+ "RefinedWebModel": ["query_key_value"],
+ "RefinedWeb": ["query_key_value"],
+ "falcon": ["query_key_value"],
+ "btlm": ["c_proj", "c_attn"],
+ "codegen": ["qkv_proj"],
+ "mistral": ["q_proj", "v_proj"],
+ "stablelm": ["q_proj", "v_proj"],
+ "phi": ["Wqkv", "out_proj", "fc1", "fc2"],
+}
diff --git a/geniusrise_text/classification/__init__.py b/geniusrise_text/classification/__init__.py
new file mode 100644
index 0000000..98390d1
--- /dev/null
+++ b/geniusrise_text/classification/__init__.py
@@ -0,0 +1,18 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .api import TextClassificationAPI
+from .bulk import TextClassificationBulk
+from .fine_tune import TextClassificationFineTuner
diff --git a/geniusrise_text/classification/api.py b/geniusrise_text/classification/api.py
new file mode 100644
index 0000000..46f5348
--- /dev/null
+++ b/geniusrise_text/classification/api.py
@@ -0,0 +1,183 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import Any, Dict
+
+import cherrypy
+import numpy as np
+import torch
+from geniusrise import BatchInput, BatchOutput, State
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
+
+from geniusrise_text.base import TextAPI
+
+log = logging.getLogger(__file__)
+
+
+class TextClassificationAPI(TextAPI):
+ r"""
+ TextClassificationAPI leveraging Hugging Face's transformers for text classification tasks.
+ This API provides an interface to classify text into various categories like sentiment, topic, intent, etc.
+
+ Attributes:
+ model (AutoModelForSequenceClassification): A Hugging Face model for sequence classification.
+ tokenizer (AutoTokenizer): A tokenizer for preprocessing text.
+ hf_pipeline (Pipeline): A Hugging Face pipeline for text classification.
+
+ Methods:
+ classify(self): Classifies text using the model and tokenizer.
+ classification_pipeline(self): Classifies text using the Hugging Face pipeline.
+ initialize_pipeline(self): Lazy initialization of the classification pipeline.
+
+ Example CLI Usage:
+ ```bash
+ genius TextClassificationAPI rise \
+ batch \
+ --input_folder ./input \
+ batch \
+ --output_folder ./output \
+ none \
+ --id cardiffnlp/twitter-roberta-base-hate-multiclass-latest-lol \
+ listen \
+ --args \
+ model_name="cardiffnlp/twitter-roberta-base-hate-multiclass-latest" \
+ model_class="AutoModelForSequenceClassification" \
+ tokenizer_class="AutoTokenizer" \
+ use_cuda=True \
+ precision="float" \
+ quantization=0 \
+ device_map="cuda:0" \
+ max_memory=None \
+ torchscript=False \
+ endpoint="*" \
+ port=3000 \
+ cors_domain="http://localhost:3000" \
+ username="user" \
+ password="password"
+ ```
+ """
+
+ def __init__(
+ self,
+ input: BatchInput,
+ output: BatchOutput,
+ state: State,
+ **kwargs,
+ ) -> None:
+ """
+ Initializes the TextClassificationAPI with the necessary configurations for input, output, and state management.
+
+ Args:
+ input (BatchInput): Configuration for the input data.
+ output (BatchOutput): Configuration for the output data.
+ state (State): State management for the API.
+ **kwargs: Additional keyword arguments for extended functionality.
+ """
+ super().__init__(input=input, output=output, state=state)
+ log.info("Loading Hugging Face API server")
+ self.hf_pipeline = None
+
+ @cherrypy.expose
+ @cherrypy.tools.json_in()
+ @cherrypy.tools.json_out()
+ @cherrypy.tools.allow(methods=["POST"])
+ def classify(self) -> Dict[str, Any]:
+ """
+ Accepts text input and returns classification results. The method uses the model and tokenizer to classify the text
+ and provide the likelihood of each class label.
+
+ Returns:
+ Dict[str, Any]: A dictionary containing the original input text and the classification scores for each label.
+
+ Example CURL Request for text classification:
+ ```bash
+ /usr/bin/curl -X POST localhost:3000/api/v1/classify \
+ -H "Content-Type: application/json" \
+ -d '{
+ "text": "tata sons lost a major contract to its rival mahindra motors"
+ }' | jq
+ ```
+ """
+ data: Dict[str, str] = cherrypy.request.json
+ text = data.get("text", "")
+
+ inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+
+ with torch.no_grad():
+ outputs = self.model(**inputs)
+ logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
+ if next(self.model.parameters()).is_cuda:
+ logits = logits.cpu()
+
+ # Handling a single number output
+ if logits.numel() == 1:
+ logits = outputs.logits.cpu().detach().numpy()
+ scores = 1 / (1 + np.exp(-logits)).flatten()
+ return {"input": text, "label_scores": scores.tolist()}
+ else:
+ softmax = torch.nn.functional.softmax(logits, dim=-1)
+ scores = softmax.numpy().tolist()
+
+ id_to_label = dict(enumerate(self.model.config.id2label.values())) # type: ignore
+ label_scores = {id_to_label[label_id]: score for label_id, score in enumerate(scores[0])}
+
+ return {"input": text, "label_scores": label_scores}
+
+ def initialize_pipeline(self):
+ """
+ Lazy initialization of the Hugging Face pipeline for classification.
+ """
+ if not self.hf_pipeline:
+ model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
+ tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+ if self.use_cuda:
+ model.cuda()
+ self.hf_pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
+
+ @cherrypy.expose
+ @cherrypy.tools.json_in()
+ @cherrypy.tools.json_out()
+ @cherrypy.tools.allow(methods=["POST"])
+ def classification_pipeline(self) -> Dict[str, Any]:
+ """
+ Accepts text input and returns classification results using the Hugging Face pipeline.
+
+ This method uses the Hugging Face pipeline for efficient and robust text classification. It's suitable for various
+ classification tasks such as sentiment analysis, topic classification, and intent recognition.
+
+ Args:
+ None - Expects input through the POST request's JSON body.
+
+ Returns:
+ Dict[str, Any]: A dictionary containing the original input text and the classification results.
+
+ Example CURL Request for text classification:
+ ```bash
+ /usr/bin/curl -X POST localhost:3000/api/v1/classification_pipeline \
+ -H "Content-Type: application/json" \
+ -d '{"text": "The movie was fantastic, with great acting and plot."}' | jq
+ ```
+ """
+ data: Dict[str, str] = cherrypy.request.json
+ text = data.get("text", "")
+
+ self.initialize_pipeline()
+ result = self.hf_pipeline(text) # type: ignore
+
+ return {"input": text, "result": result}
diff --git a/geniusrise_text/classification/api.yml b/geniusrise_text/classification/api.yml
new file mode 100644
index 0000000..f243808
--- /dev/null
+++ b/geniusrise_text/classification/api.yml
@@ -0,0 +1,88 @@
+openapi: 3.0.0
+info:
+ title: GeniusRise Text Processing API
+ description: API for text generation and classification using Hugging Face models.
+ version: "1.0"
+servers:
+ - url: http://localhost:3000/api/v1
+ description: Development server
+paths:
+
+ /classify:
+ post:
+ summary: Classifies text into categories
+ operationId: classifyText
+ tags:
+ - Text Classification
+ requestBody:
+ required: true
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ text:
+ type: string
+ description: The text to classify
+ required:
+ - text
+ responses:
+ 200:
+ description: Successfully classified text
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ input:
+ type: string
+ label_scores:
+ type: object
+ additionalProperties:
+ type: number
+ 400:
+ description: Invalid request
+ 500:
+ description: Error during text classification
+
+ /classification_pipeline:
+ post:
+ summary: Classifies text using the Hugging Face pipeline
+ operationId: classifyTextPipeline
+ tags:
+ - Text Classification
+ requestBody:
+ required: true
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ text:
+ type: string
+ description: The text to classify
+ required:
+ - text
+ responses:
+ 200:
+ description: Successfully classified text using pipeline
+ content:
+ application/json:
+ schema:
+ type: object
+ properties:
+ input:
+ type: string
+ result:
+ type: array
+ items:
+ type: object
+ properties:
+ label:
+ type: string
+ score:
+ type: number
+ 400:
+ description: Invalid request
+ 500:
+ description: Error during text classification
diff --git a/geniusrise_text/classification/bulk.py b/geniusrise_text/classification/bulk.py
new file mode 100644
index 0000000..0152b47
--- /dev/null
+++ b/geniusrise_text/classification/bulk.py
@@ -0,0 +1,363 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import json
+import os
+import sqlite3
+import uuid
+import xml.etree.ElementTree as ET
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+import torch
+import yaml # type: ignore
+from datasets import Dataset, load_from_disk
+from geniusrise import BatchInput, BatchOutput, State
+from pyarrow import feather
+from pyarrow import parquet as pq
+
+from geniusrise_text.base import TextBulk
+
+
+class TextClassificationBulk(TextBulk):
+ r"""
+ TextClassificationBulk is designed to handle bulk text classification tasks using Hugging Face models efficiently and
+ effectively. It allows for processing large datasets, utilizing state-of-the-art machine learning models to provide
+ accurate classification of text data into predefined labels.
+
+ Args:
+ input (BatchInput): Configuration and data inputs for the batch process.
+ output (BatchOutput): Configurations for output data handling.
+ state (State): State management for the classification task.
+ **kwargs: Arbitrary keyword arguments for extended configurations.
+
+ Example CLI Usage:
+ ```bash
+ genius TextClassificationBulk rise \
+ batch \
+ --input_folder ./input \
+ batch \
+ --output_folder ./output \
+ none \
+ --id cardiffnlp/twitter-roberta-base-hate-multiclass-latest-lol \
+ classify \
+ --args \
+ model_name="cardiffnlp/twitter-roberta-base-hate-multiclass-latest" \
+ model_class="AutoModelForSequenceClassification" \
+ tokenizer_class="AutoTokenizer" \
+ use_cuda=True \
+ precision="bfloat16" \
+ quantization=0 \
+ device_map="auto" \
+ max_memory=None \
+ torchscript=False
+ ```
+ """
+
+ def __init__(self, input: BatchInput, output: BatchOutput, state: State, **kwargs) -> None:
+ """
+ Initializes the TextClassificationBulk class with input, output, and state configurations.
+
+ Args:
+ input (BatchInput): Configuration for the input data.
+ output (BatchOutput): Configuration for the output data.
+ state (State): State management for the classification task.
+ **kwargs: Additional keyword arguments for extended functionality.
+ """
+ super().__init__(input, output, state, **kwargs)
+
+ def load_dataset(self, dataset_path: str, max_length: int = 512, **kwargs) -> Optional[Dataset]:
+ r"""
+ Load a classification dataset from a directory.
+
+ Args:
+ dataset_path (str): The path to the dataset directory.
+ max_length (int, optional): The maximum length for tokenization. Defaults to 512.
+
+ Returns:
+ Dataset: The loaded dataset.
+
+ Raises:
+ Exception: If there was an error loading the dataset.
+
+ ## Supported Data Formats and Structures:
+
+ ### JSONL
+ Each line is a JSON object representing an example.
+ ```json
+ {"text": "The text content"}
+ ```
+
+ ### CSV
+ Should contain 'text' columns.
+ ```csv
+ text
+ "The text content"
+ ```
+
+ ### Parquet
+ Should contain 'text' columns.
+
+ ### JSON
+ An array of dictionaries with 'text' keys.
+ ```json
+ [{"text": "The text content"}]
+ ```
+
+ ### XML
+ Each 'record' element should contain 'text' child elements.
+ ```xml
+
+ The text content
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'text' keys.
+ ```yaml
+ - text: "The text content"
+ ```
+
+ ### TSV
+ Should contain 'text' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'text' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'text' columns.
+
+ ### Feather
+ Should contain 'text' columns.
+ """
+ self.max_length = max_length
+
+ self.label_to_id = self.model.config.label2id if self.model and self.model.config.label2id else {} # type: ignore
+
+ try:
+ self.log.info(f"Loading dataset from {dataset_path}")
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ # Load dataset saved by Hugging Face datasets library
+ return load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ text = record.find("text").text # type: ignore
+ data.append({"text": text})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT text FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ data = [fn(d) for d in data]
+ else:
+ data = data
+
+ return Dataset.from_pandas(pd.DataFrame(data))
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def classify(
+ self,
+ model_name: str,
+ model_class: str = "AutoModelForSequenceClassification",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Perform bulk classification using the specified model and tokenizer. This method handles the entire classification
+ process including loading the model, processing input data, predicting classifications, and saving the results.
+
+ Args:
+ model_name (str): Name or path of the model.
+ model_class (str): Class name of the model (default "AutoModelForSequenceClassification").
+ tokenizer_class (str): Class name of the tokenizer (default "AutoTokenizer").
+ use_cuda (bool): Whether to use CUDA for model inference (default False).
+ precision (str): Precision for model computation (default "float").
+ quantization (int): Level of quantization for optimizing model size and speed (default 0).
+ device_map (str | Dict | None): Specific device to use for computation (default "auto").
+ max_memory (Dict): Maximum memory configuration for devices.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool): Whether to enable AWQ optimization (default False).
+ flash_attention (bool): Whether to use flash attention optimization (default False).
+ batch_size (int): Number of classifications to process simultaneously (default 32).
+ **kwargs: Arbitrary keyword arguments for model and generation configurations.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["text"]
+
+ # Process data in batches
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+ inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
+
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+
+ predictions = self.model(**inputs)
+ predictions = predictions[0] if isinstance(predictions, tuple) else predictions.logits
+ predictions = torch.argmax(predictions, dim=-1).cpu().numpy()
+
+ self._save_predictions(predictions, batch, output_path, i)
+ self.done()
+
+ def _save_predictions(
+ self,
+ predictions: torch.Tensor,
+ input_batch: List[str],
+ output_path: str,
+ batch_idx: int,
+ ) -> None:
+ """
+ Saves the classification predictions to a specified output path. This method is called internally by the classify method
+ to persist the classification results.
+
+ Args:
+ predictions (torch.Tensor): Tensor of label indices predicted by the model.
+ input_batch (List[str]): List of original texts that were classified.
+ output_path (str): Path to save the classification results.
+ batch_idx (int): Index of the current batch (for naming files).
+ """
+ id_to_label = dict(enumerate(self.model.config.id2label.values())) # type: ignore
+ label_predictions = [id_to_label[label_id] for label_id in predictions.tolist()]
+
+ # Prepare data for saving
+ data_to_save = [
+ {"input": input_text, "prediction": label} for input_text, label in zip(input_batch, label_predictions)
+ ]
+ with open(
+ os.path.join(output_path, f"predictions-{batch_idx}-{str(uuid.uuid4())}.json"),
+ "w",
+ ) as f:
+ json.dump(data_to_save, f)
diff --git a/geniusrise_text/classification/eval.py b/geniusrise_text/classification/eval.py
new file mode 100644
index 0000000..1ecde3b
--- /dev/null
+++ b/geniusrise_text/classification/eval.py
@@ -0,0 +1,412 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import sqlite3
+import uuid
+import xml.etree.ElementTree as ET
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+import torch
+import yaml # type: ignore
+from datasets import Dataset, load_from_disk, load_dataset, load_metric
+from geniusrise import BatchInput, BatchOutput, State
+from pyarrow import feather
+from pyarrow import parquet as pq
+
+from geniusrise_text.base import TextBulk
+
+
+class TextClassificationEval(TextBulk):
+ r"""
+ TextClassificationEval extends TextBulk to support evaluation of text classification models on large datasets. It facilitates
+ processing of datasets, model inference, and computation of evaluation metrics such as accuracy, precision, recall, and F1 score.
+
+ Args:
+ input (BatchInput): Configuration and data inputs for the batch process.
+ output (BatchOutput): Configurations for output data handling.
+ state (State): State management for the classification task.
+ **kwargs: Arbitrary keyword arguments for extended configurations.
+
+ Example CLI Usage:
+ ```bash
+ genius TextClassificationEval rise \
+ batch \
+ --input_folder ./input \
+ batch \
+ --output_folder ./output \
+ none \
+ --id cardiffnlp/twitter-roberta-base-hate-multiclass-latest-lol \
+ classify \
+ --args \
+ model_name="cardiffnlp/twitter-roberta-base-hate-multiclass-latest" \
+ model_class="AutoModelForSequenceClassification" \
+ tokenizer_class="AutoTokenizer" \
+ use_cuda=True \
+ precision="bfloat16" \
+ quantization=0 \
+ device_map="auto" \
+ max_memory=None \
+ torchscript=False
+ ```
+ """
+
+ def __init__(self, input: BatchInput, output: BatchOutput, state: State, **kwargs) -> None:
+ """
+ Initializes the TextClassificationEval class with configurations for input, output, state, and evaluation settings.
+
+ Args:
+ input (BatchInput): Configuration for the input data.
+ output (BatchOutput): Configuration for the output data.
+ state (State): State management for the classification task.
+ **kwargs: Additional keyword arguments for extended functionality.
+ """
+ super().__init__(input, output, state, **kwargs)
+
+ def load_dataset(self, dataset_path: str, max_length: int = 512, **kwargs) -> Optional[Dataset]:
+ r"""
+ Load a classification dataset from a directory.
+
+ Args:
+ dataset_path (str): The path to the dataset directory.
+ max_length (int, optional): The maximum length for tokenization. Defaults to 512.
+
+ Returns:
+ Dataset: The loaded dataset.
+
+ Raises:
+ Exception: If there was an error loading the dataset.
+
+ ## Supported Data Formats and Structures:
+
+ ### JSONL
+ Each line is a JSON object representing an example.
+ ```json
+ {"text": "The text content", "label": "The label"}
+ ```
+
+ ### CSV
+ Should contain 'text' and 'label' columns.
+ ```csv
+ text,label
+ "The text content","The label"
+ ```
+
+ ### Parquet
+ Should contain 'text' and 'label' columns.
+
+ ### JSON
+ An array of dictionaries with 'text' and 'label' keys.
+ ```json
+ [{"text": "The text content", "label": "The label"}]
+ ```
+
+ ### XML
+ Each 'record' element should contain 'text' and 'label' child elements.
+ ```xml
+
+ The text content
+ The label
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'text' and 'label' keys.
+ ```yaml
+ - text: "The text content"
+ label: "The label"
+ ```
+
+ ### TSV
+ Should contain 'text' and 'label' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'text' and 'label' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'text' and 'label' columns.
+
+ ### Feather
+ Should contain 'text' and 'label' columns.
+ """
+ self.max_length = max_length
+
+ self.label_to_id = self.model.config.label2id if self.model and self.model.config.label2id else {} # type: ignore
+
+ try:
+ self.log.info(f"Loading dataset from {dataset_path}")
+ if self.use_huggingface_dataset:
+ dataset = load_dataset(self.huggingface_dataset)
+ elif os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+
+ dataset = load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in os.listdir(dataset_path):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ text = record.find("text").text # type: ignore
+ label = record.find("label").text # type: ignore
+ data.append({"text": text, "label": label})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT text, label FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ dataset = Dataset.from_pandas(pd.DataFrame(data))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ dataset = dataset.map(fn)
+ else:
+ dataset = dataset
+
+ return dataset
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def evaluate(
+ self,
+ model_name: str,
+ model_class: str = "AutoModelForSequenceClassification",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ notification_email: Optional[str] = None,
+ use_huggingface_dataset: bool = False,
+ huggingface_dataset: str = "",
+ **kwargs: Any,
+ ) -> None:
+ """
+ Evaluates the model on the loaded dataset, calculates evaluation metrics, and saves both predictions and metrics.
+
+ Args:
+ model_name (str): Name or path of the model.
+ model_class (str): Class name of the model (default "AutoModelForSequenceClassification").
+ tokenizer_class (str): Class name of the tokenizer (default "AutoTokenizer").
+ use_cuda (bool): Whether to use CUDA for model inference (default False).
+ precision (str): Precision for model computation (default "float").
+ quantization (int): Level of quantization for optimizing model size and speed (default 0).
+ device_map (str | Dict | None): Specific device to use for computation (default "auto").
+ max_memory (Dict): Maximum memory configuration for devices.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool): Whether to enable AWQ optimization (default False).
+ flash_attention (bool): Whether to use flash attention optimization (default False).
+ batch_size (int): Number of classifications to process simultaneously (default 32).
+ use_huggingface_dataset (bool, optional): Whether to load a dataset from huggingface hub.
+ huggingface_dataset (str, optional): The huggingface dataset to use.
+ **kwargs: Arbitrary keyword arguments for model and generation configurations.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.compile = compile
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.use_huggingface_dataset = use_huggingface_dataset
+ self.huggingface_dataset = huggingface_dataset
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["text"]
+
+ # Ensure metrics are available
+ accuracy_metric = load_metric("accuracy")
+ precision_metric = load_metric("precision")
+ recall_metric = load_metric("recall")
+ f1_metric = load_metric("f1")
+
+ all_predictions = []
+ all_true_labels = []
+
+ # Loop through the dataset in batches
+ for i in range(0, len(dataset), batch_size):
+ batch_texts = dataset[i : i + batch_size]["text"]
+ batch_labels = dataset[i : i + batch_size]["labels"]
+ inputs = self.tokenizer(
+ batch_texts, padding=True, truncation=True, max_length=self.max_length, return_tensors="pt"
+ )
+
+ if self.use_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+ batch_labels = torch.tensor(batch_labels).cuda()
+
+ with torch.no_grad():
+ outputs = self.model(**inputs)
+ predictions = torch.argmax(outputs.logits, dim=-1)
+
+ all_predictions.extend(predictions.cpu().numpy())
+ all_true_labels.extend(batch_labels.cpu().numpy())
+
+ # Compute overall metrics
+ # fmt: off
+ overall_accuracy = accuracy_metric.compute(predictions=all_predictions, references=all_true_labels)["accuracy"]
+ overall_precision = precision_metric.compute(predictions=all_predictions, references=all_true_labels, average="macro")["precision"]
+ overall_recall = recall_metric.compute(predictions=all_predictions, references=all_true_labels, average="macro")["recall"]
+ overall_f1 = f1_metric.compute(predictions=all_predictions, references=all_true_labels, average="macro")["f1"]
+ # fmt: on
+
+ overall_evaluation_metrics = {
+ "accuracy": overall_accuracy,
+ "precision": overall_precision,
+ "recall": overall_recall,
+ "f1": overall_f1,
+ }
+
+ # Save predictions and evaluation metrics
+ self._save_predictions(
+ all_predictions, dataset["text"], all_true_labels, self.output.output_folder, overall_evaluation_metrics
+ )
+
+ self.done()
+
+ def _save_predictions(
+ self,
+ predictions: List[int],
+ input_texts: List[str],
+ true_labels: List[int],
+ output_path: str,
+ evaluation_metrics: Dict[str, float],
+ ) -> None:
+ """
+ Saves the classification predictions and evaluation metrics to a specified output path.
+
+ Args:
+ predictions (List[int]): List of label indices predicted by the model.
+ input_texts (List[str]): List of original texts that were classified.
+ true_labels (List[int]): List of true label indices.
+ output_path (str): Path to save the classification results and metrics.
+ evaluation_metrics (Dict[str, float]): Dictionary of evaluation metrics.
+ """
+ # Prepare data for saving
+ data_to_save = [
+ {"input": input_text, "prediction": prediction, "true_label": true_label}
+ for input_text, prediction, true_label in zip(input_texts, predictions, true_labels)
+ ]
+
+ # Save predictions to a JSON file
+ predictions_file_path = os.path.join(output_path, f"predictions-{str(uuid.uuid4())}.json")
+ with open(predictions_file_path, "w") as f:
+ json.dump(data_to_save, f, ensure_ascii=False, indent=4)
+
+ # Save evaluation metrics to a JSON file
+ metrics_file_path = os.path.join(output_path, "evaluation_metrics.json")
+ with open(metrics_file_path, "w") as f:
+ json.dump(evaluation_metrics, f, ensure_ascii=False, indent=4)
diff --git a/geniusrise_text/classification/fine_tune.py b/geniusrise_text/classification/fine_tune.py
new file mode 100644
index 0000000..33bcf51
--- /dev/null
+++ b/geniusrise_text/classification/fine_tune.py
@@ -0,0 +1,290 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import sqlite3
+import xml.etree.ElementTree as ET
+from typing import Dict, Optional, Union
+
+import numpy as np
+import pandas as pd
+import yaml # type: ignore
+from datasets import Dataset, load_dataset, load_from_disk
+from pyarrow import feather
+from pyarrow import parquet as pq
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+from transformers import DataCollatorWithPadding, EvalPrediction
+
+from geniusrise_text.base import TextFineTuner
+
+
+class TextClassificationFineTuner(TextFineTuner):
+ r"""
+ A bolt for fine-tuning Hugging Face models for text classification tasks.
+
+ This class extends the `TextFineTuner` and specializes in fine-tuning models for text classification.
+ It provides additional functionalities for loading and preprocessing text classification datasets in various formats.
+
+ Args:
+ input (BatchInput): The batch input data.
+ output (OutputConfig): The output data.
+ state (State): The state manager.
+
+ CLI Usage:
+
+ ```bash
+ genius TextClassificationFineTuner rise \
+ batch \
+ --input_folder ./input \
+ batch \
+ --output_folder ./output \
+ none \
+ --id cardiffnlp/twitter-roberta-base-hate-multiclass-latest-lol \
+ fine_tune \
+ --args \
+ model_name=my_model \
+ tokenizer_name=my_tokenizer \
+ num_train_epochs=3 \
+ per_device_train_batch_size=8 \
+ data_max_length=512
+ ```
+ """
+
+ def load_dataset(self, dataset_path: str, max_length: int = 512, **kwargs) -> Optional[Dataset]:
+ r"""
+ Load a classification dataset from a directory.
+
+ Args:
+ dataset_path (str): The path to the dataset directory.
+ max_length (int, optional): The maximum length for tokenization. Defaults to 512.
+
+ Returns:
+ Dataset: The loaded dataset.
+
+ Raises:
+ Exception: If there was an error loading the dataset.
+
+ ## Supported Data Formats and Structures:
+
+ ### JSONL
+ Each line is a JSON object representing an example.
+ ```json
+ {"text": "The text content", "label": "The label"}
+ ```
+
+ ### CSV
+ Should contain 'text' and 'label' columns.
+ ```csv
+ text,label
+ "The text content","The label"
+ ```
+
+ ### Parquet
+ Should contain 'text' and 'label' columns.
+
+ ### JSON
+ An array of dictionaries with 'text' and 'label' keys.
+ ```json
+ [{"text": "The text content", "label": "The label"}]
+ ```
+
+ ### XML
+ Each 'record' element should contain 'text' and 'label' child elements.
+ ```xml
+
+ The text content
+ The label
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'text' and 'label' keys.
+ ```yaml
+ - text: "The text content"
+ label: "The label"
+ ```
+
+ ### TSV
+ Should contain 'text' and 'label' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'text' and 'label' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'text' and 'label' columns.
+
+ ### Feather
+ Should contain 'text' and 'label' columns.
+ """
+
+ self.data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer, max_length=max_length)
+ self.max_length = max_length
+
+ # self.label_to_id = self.model.config.label2id if self.model and self.model.config.label2id else {} # type: ignore
+ self.label_to_id = {}
+
+ def tokenize_function(examples):
+ tokenized_data = self.tokenizer(
+ examples["text"],
+ padding="max_length",
+ truncation=True,
+ max_length=self.max_length,
+ )
+
+ labels = [x for x in list(set(examples["label"]))]
+ all_labels = [l for l in examples["label"]]
+
+ unknown_labels = [label for label in labels if label not in self.label_to_id]
+
+ self.label_to_id = {
+ **self.label_to_id,
+ **{x: i for i, x in enumerate(unknown_labels)},
+ }
+
+ tokenized_data["label"] = [self.label_to_id[label] for label in all_labels]
+ return tokenized_data
+
+ try:
+ self.log.info(f"Loading dataset from {dataset_path}")
+ if self.use_huggingface_dataset:
+ dataset = load_dataset(self.huggingface_dataset)
+ elif os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ # Load dataset saved by Hugging Face datasets library
+ dataset = load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in os.listdir(dataset_path):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ text = record.find("text").text # type: ignore
+ label = record.find("label").text # type: ignore
+ data.append({"text": text, "label": label})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT text, label FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ dataset = Dataset.from_pandas(pd.DataFrame(data))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ dataset = dataset.map(fn)
+ else:
+ dataset = dataset
+
+ # Create label_to_id mapping and save it in model config
+ # TODO: ugly shit cause we dont know num labels before we process the data but need tokenizer to process data
+ self.label_to_id = {label: i for i, label in enumerate(set(dataset["train"]["label"]))}
+ if self.model:
+ config = self.model.config
+ config.label2id = self.label_to_id
+ config.id2label = {i: label for label, i in self.label_to_id.items()}
+ config.num_labels = len(self.label_to_id.keys())
+ self.config = config
+
+ self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ device_map=self.device_map,
+ precision=self.precision,
+ quantization=self.quantization,
+ lora_config=self.lora_config,
+ use_accelerate=self.use_accelerate,
+ accelerate_no_split_module_classes=self.accelerate_no_split_module_classes,
+ **self.model_kwargs,
+ )
+ if self.tokenizer and not self.tokenizer.pad_token:
+ self.tokenizer.pad_token = self.tokenizer.eos_token
+ self.model.config.pad_token_id = self.tokenizer.eos_token_id
+
+ self.log.info(self.model.config)
+
+ return dataset.map(tokenize_function, batched=True)
+
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def compute_metrics(self, eval_pred: EvalPrediction) -> Union[Optional[Dict[str, float]], Dict[str, float]]:
+ """
+ Compute metrics for evaluation. This class implements a simple classification evaluation,
+ tasks should ideally override this.
+
+ Args:
+ eval_pred (EvalPrediction): The evaluation predictions.
+
+ Returns:
+ dict: The computed metrics.
+ """
+ predictions, labels = eval_pred
+ predictions = predictions[0] if isinstance(predictions, tuple) else predictions
+ labels = labels[0] if isinstance(labels, tuple) else labels
+ predictions = np.argmax(predictions, axis=1)
+
+ is_binary = len(self.label_to_id.keys()) == 2
+ average_type = "binary" if is_binary else "weighted"
+
+ precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average=average_type)
+
+ return {
+ "accuracy": accuracy_score(labels, predictions),
+ "precision": precision,
+ "recall": recall,
+ "f1": f1,
+ }
diff --git a/geniusrise_text/classification/tests/__init__.py b/geniusrise_text/classification/tests/__init__.py
new file mode 100644
index 0000000..bffa8ca
--- /dev/null
+++ b/geniusrise_text/classification/tests/__init__.py
@@ -0,0 +1,14 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/geniusrise_text/classification/tests/test_bulk.py b/geniusrise_text/classification/tests/test_bulk.py
new file mode 100644
index 0000000..3063aaf
--- /dev/null
+++ b/geniusrise_text/classification/tests/test_bulk.py
@@ -0,0 +1,186 @@
+# 🧠 Geniusrise
+# Copyright (C) 2023 geniusrise.ai
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import glob
+import json
+import os
+import sqlite3
+import tempfile
+import xml.etree.ElementTree as ET
+
+import pandas as pd
+import pytest
+import yaml # type: ignore
+from datasets import Dataset
+from geniusrise.core import BatchInput, BatchOutput, InMemoryState
+from pyarrow import feather
+from pyarrow import parquet as pq
+
+from geniusrise_text.classification.bulk import TextClassificationBulk
+
+MODELS_TO_TEST = {
+ # fmt: off
+ "cardiffnlp/twitter-roberta-base-hate-multiclass-latest": ["sexism", "racism", "disability", "sexual_orientation", "religion", "other", "not_hate"],
+ "cardiffnlp/twitter-roberta-base-hate-latest": ["NOT-HATE", "HATE"],
+ "cardiffnlp/twitter-roberta-base-offensive": ["non-offensive", "offensive"],
+ "cardiffnlp/twitter-xlm-roberta-base-sentiment": ["positive", "neutral", "negative"],
+ "cardiffnlp/twitter-roberta-base-emotion": ["joy", "optimism", "anger", "sadness"],
+ "cardiffnlp/twitter-roberta-base-irony": ["non_irony", "irony"],
+ "cardiffnlp/twitter-xlm-roberta-base-sentiment-multilingual": ["positive", "neutral", "negative"],
+ "tomh/toxigen_roberta": ["LABEL_0", "LABEL_1"],
+ "cointegrated/rubert-tiny-toxicity": ["non-toxic", "insult", "obscenity", "threat", "dangerous"],
+ "michellejieli/NSFW_text_classifier": ["SFW", "NSFW"],
+ "bvanaken/clinical-assertion-negation-bert": ["ABSENT", "PRESENT", "POSSIBLE"],
+ "bucketresearch/politicalBiasBERT": ["LEFT", "CENTER", "RIGHT"],
+ "soleimanian/financial-roberta-large-sentiment": ["neutral", "negative", "positive"],
+ "jpwahle/longformer-base-plagiarism-detection": ["ORIGINAL", "PLAGIARISM"],
+ "mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis": ["negative", "neutral", "positive"],
+ "lxyuan/distilbert-base-multilingual-cased-sentiments-student": ["negative", "neutral", "positive"],
+ "Sigma/financial-sentiment-analysis": ["LABEL_0", "LABEL_1", "LABEL_2"],
+ "SamLowe/roberta-base-go_emotions": ["disappointment", "sadness", "annoyance", "neutral", "disapproval", "realization",
+ "nervousness", "approval", "joy", "anger", "embarrassment", "caring", "remorse",
+ "disgust", "grief", "confusion", "relief", "desire", "admiration", "optimism",
+ "fear", "love", "excitement", "curiosity", "amusement", "surprise", "gratitude", "pride"],
+ "cardiffnlp/tweet-topic-21-multi": ["sports", "news_&_social_concern", "fitness_&_health", "youth_&_student_life", "learning_&_educational",
+ "science_&_technology", "celebrity_&_pop_culture", "travel_&_adventure", "diaries_&_daily_life",
+ "food_&_dining", "gaming", "business_&_entrepreneurs", "family", "relationships", "fashion_&_style",
+ "music", "film_tv_&_video", "other_hobbies", "arts_&_culture"],
+ "padmajabfrl/Gender-Classification": ["Female", "Male"],
+ "ProsusAI/finbert": ["positive", "neutral", "negative"],
+ "yiyanghkust/finbert-tone": ["Positive", "Neutral", "Negative"],
+ "wajidlinux99/gibberish-text-detector": ["clean", "mild gibberish", "word salad", "noise"],
+ "cnut1648/biolinkbert-large-mnli-snli": ["entailment", "neutral", "contradiction"],
+ # fmt: on
+}
+
+
+# Helper function to create synthetic data in different formats
+def create_dataset_in_format(directory, ext):
+ os.makedirs(directory, exist_ok=True)
+ data = [{"text": f"text_{i}"} for i in range(10)]
+ df = pd.DataFrame(data)
+
+ if ext == "huggingface":
+ dataset = Dataset.from_pandas(df)
+ dataset.save_to_disk(directory)
+ elif ext == "csv":
+ df.to_csv(os.path.join(directory, "data.csv"), index=False)
+ elif ext == "jsonl":
+ with open(os.path.join(directory, "data.jsonl"), "w") as f:
+ for item in data:
+ f.write(json.dumps(item) + "\n")
+ elif ext == "parquet":
+ pq.write_table(feather.Table.from_pandas(df), os.path.join(directory, "data.parquet"))
+ elif ext == "json":
+ with open(os.path.join(directory, "data.json"), "w") as f:
+ json.dump(data, f)
+ elif ext == "xml":
+ root = ET.Element("root")
+ for item in data:
+ record = ET.SubElement(root, "record")
+ ET.SubElement(record, "text").text = item["text"]
+ tree = ET.ElementTree(root)
+ tree.write(os.path.join(directory, "data.xml"))
+ elif ext == "yaml":
+ with open(os.path.join(directory, "data.yaml"), "w") as f:
+ yaml.dump(data, f)
+ elif ext == "tsv":
+ df.to_csv(os.path.join(directory, "data.tsv"), index=False, sep="\t")
+ elif ext == "xlsx":
+ df.to_excel(os.path.join(directory, "data.xlsx"), index=False)
+ elif ext == "db":
+ conn = sqlite3.connect(os.path.join(directory, "data.db"))
+ df.to_sql("dataset_table", conn, if_exists="replace", index=False)
+ conn.close()
+ elif ext == "feather":
+ feather.write_feather(df, os.path.join(directory, "data.feather"))
+
+
+# Fixture for models
+@pytest.fixture(params=MODELS_TO_TEST.items())
+def model(request):
+ return request.param
+
+
+# Fixtures for each file type
+@pytest.fixture(
+ params=[
+ "huggingface",
+ "csv",
+ "jsonl",
+ "parquet",
+ "json",
+ "xml",
+ "yaml",
+ "tsv",
+ "xlsx",
+ "db",
+ "feather",
+ ]
+)
+def dataset_file(request, tmpdir):
+ ext = request.param
+ create_dataset_in_format(tmpdir, ext)
+ return tmpdir, ext
+
+
+@pytest.fixture
+def classification_bolt():
+ input_dir = tempfile.mkdtemp()
+ output_dir = tempfile.mkdtemp()
+ input = BatchInput(input_dir, "geniusrise-test", "test-🤗-input")
+ output = BatchOutput(output_dir, "geniusrise-test", "test-🤗-output")
+ state = InMemoryState()
+ klass = TextClassificationBulk(
+ input=input,
+ output=output,
+ state=state,
+ )
+
+ return klass
+
+
+def test_classify(classification_bolt, dataset_file, model):
+ tmpdir, ext = dataset_file
+ classification_bolt.input.input_folder = tmpdir
+
+ model_name, labels = model
+ tokenizer_name = model_name
+ model_class = "AutoModelForSequenceClassification"
+ tokenizer_class = "AutoTokenizer"
+
+ # Classify
+ classification_bolt.classify(
+ model_name=model_name,
+ tokenizer_name=tokenizer_name,
+ model_class=model_class,
+ tokenizer_class=tokenizer_class,
+ device_map="cuda:0",
+ num_train_epochs=2,
+ per_device_batch_size=2,
+ precision="float16",
+ )
+ # Check output
+ files = glob.glob(f"{classification_bolt.output.output_folder}/predictions-*.json")
+ assert len(files) > 0
+
+ # Check one of the output files to ensure it contains the predictions and input data
+ with open(files[0], "r") as f:
+ results = json.load(f)
+ assert len(results) > 0
+ for result in results:
+ assert "input" in result
+ assert "prediction" in result
+ assert result["prediction"] in MODELS_TO_TEST[model_name]
diff --git a/tests/test_classification.py b/geniusrise_text/classification/tests/test_fine_tune.py
similarity index 58%
rename from tests/test_classification.py
rename to geniusrise_text/classification/tests/test_fine_tune.py
index 8b335e4..9b401a4 100644
--- a/tests/test_classification.py
+++ b/geniusrise_text/classification/tests/test_fine_tune.py
@@ -1,18 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ The instruction
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'instruction' keys.
+ ```yaml
+ - instruction: "The instruction"
+ ```
+
+ ### TSV
+ Should contain 'instruction' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'instruction' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'instruction' columns.
+
+ ### Feather
+ Should contain 'instruction' columns.
+ """
+ try:
+ self.log.info(f"Loading dataset from {dataset_path}")
+ self.max_length = max_length
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ # Load dataset saved by Hugging Face datasets library
+ return load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ instruction = record.find("instruction").text # type: ignore
+ data.append({"instruction": instruction})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT instruction FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ data = [fn(d) for d in data]
+ else:
+ data = data
+
+ dataset = Dataset.from_pandas(pd.DataFrame(data))
+ return dataset
+ except Exception as e:
+ self.log.error(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def perform(
+ self,
+ model_name: str,
+ model_class: str = "AutoModelForCausalLM",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ decoding_strategy: str = "generate",
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Performs text generation in bulk using a specified instruction-tuned model. This method handles the entire
+ process, including model loading, prompt processing, text generation, and saving the results.
+
+ Args:
+ model_name (str): The name or path of the instruction-tuned model.
+ model_class (str, optional): The class of the language model. Defaults to "AutoModelForCausalLM".
+ tokenizer_class (str, optional): The class of the tokenizer. Defaults to "AutoTokenizer".
+ use_cuda (bool, optional): Whether to use CUDA for model inference. Defaults to False.
+ precision (str, optional): Precision for model computation. Defaults to "float16".
+ quantization (int, optional): Level of quantization for optimizing model size and speed. Defaults to 0.
+ device_map (str | Dict | None, optional): Specific device to use for computation. Defaults to "auto".
+ max_memory (Dict, optional): Maximum memory configuration for devices. Defaults to {0: "24GB"}.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool, optional): Whether to enable AWQ optimization. Defaults to False.
+ flash_attention (bool, optional): Whether to use flash attention optimization. Defaults to False.
+ decoding_strategy (str, optional): Strategy for decoding the completion. Defaults to "generate".
+ **kwargs: Configuration and additional arguments for text generation such as model class, tokenizer class,
+ precision, device map, and other generation-related parameters.
+
+ Note:
+ Additional arguments are passed directly to the model and tokenizer initialization and the generation method.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["instruction"]
+
+ prompts = []
+ completions = []
+
+ for _, prompt in enumerate(dataset):
+ completion = self.generate(
+ prompt=prompt,
+ decoding_strategy=decoding_strategy,
+ **generation_args,
+ )
+ completions.append(completion)
+ prompts.append(prompt)
+
+ self._save_completions(completions, prompts, output_path)
+ self.done()
+
+ def perform_vllm(
+ self,
+ model_name: str,
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ # VLLM params
+ vllm_tokenizer_mode: str = "auto",
+ vllm_download_dir: Optional[str] = None,
+ vllm_load_format: str = "auto",
+ vllm_seed: int = 42,
+ vllm_max_model_len: int = 1024,
+ vllm_enforce_eager: bool = False,
+ vllm_max_context_len_to_capture: int = 8192,
+ vllm_block_size: int = 16,
+ vllm_gpu_memory_utilization: float = 0.90,
+ vllm_swap_space: int = 4,
+ vllm_sliding_window: Optional[int] = None,
+ vllm_pipeline_parallel_size: int = 1,
+ vllm_tensor_parallel_size: int = 1,
+ vllm_worker_use_ray: bool = False,
+ vllm_max_parallel_loading_workers: Optional[int] = None,
+ vllm_disable_custom_all_reduce: bool = False,
+ vllm_max_num_batched_tokens: Optional[int] = None,
+ vllm_max_num_seqs: int = 64,
+ vllm_max_paddings: int = 512,
+ vllm_max_lora_rank: Optional[int] = None,
+ vllm_max_loras: Optional[int] = None,
+ vllm_max_cpu_loras: Optional[int] = None,
+ vllm_lora_extra_vocab_size: int = 0,
+ vllm_placement_group: Optional[dict] = None,
+ vllm_log_stats: bool = False,
+ # Generate params
+ notification_email: Optional[str] = None,
+ batch_size: int = 32,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Performs bulk text generation using the Versatile Language Learning Model (VLLM) with specified parameters
+ for fine-tuning model behavior, including quantization and parallel processing settings. This method is designed
+ to process large datasets efficiently by leveraging VLLM capabilities for generating high-quality text completions
+ based on provided prompts.
+
+ Args:
+ model_name (str): The name or path of the VLLM model to use for text generation.
+ use_cuda (bool): Flag indicating whether to use CUDA for GPU acceleration.
+ precision (str): Precision of computations, can be "float16", "bfloat16", etc.
+ quantization (int): Level of quantization for model weights, 0 for none.
+ device_map (str | Dict | None): Specific device(s) to use for model inference.
+ vllm_tokenizer_mode (str): Mode of the tokenizer ("auto", "fast", or "slow").
+ vllm_download_dir (Optional[str]): Directory to download and load the model and tokenizer.
+ vllm_load_format (str): Format to load the model, e.g., "auto", "pt".
+ vllm_seed (int): Seed for random number generation.
+ vllm_max_model_len (int): Maximum sequence length the model can handle.
+ vllm_enforce_eager (bool): Enforce eager execution instead of using optimization techniques.
+ vllm_max_context_len_to_capture (int): Maximum context length for CUDA graph capture.
+ vllm_block_size (int): Block size for caching mechanism.
+ vllm_gpu_memory_utilization (float): Fraction of GPU memory to use.
+ vllm_swap_space (int): Amount of swap space to use in GiB.
+ vllm_sliding_window (Optional[int]): Size of the sliding window for processing.
+ vllm_pipeline_parallel_size (int): Number of pipeline parallel groups.
+ vllm_tensor_parallel_size (int): Number of tensor parallel groups.
+ vllm_worker_use_ray (bool): Whether to use Ray for model workers.
+ vllm_max_parallel_loading_workers (Optional[int]): Maximum number of workers for parallel loading.
+ vllm_disable_custom_all_reduce (bool): Disable custom all-reduce kernel and fall back to NCCL.
+ vllm_max_num_batched_tokens (Optional[int]): Maximum number of tokens to be processed in a single iteration.
+ vllm_max_num_seqs (int): Maximum number of sequences to be processed in a single iteration.
+ vllm_max_paddings (int): Maximum number of paddings to be added to a batch.
+ vllm_max_lora_rank (Optional[int]): Maximum rank for LoRA adjustments.
+ vllm_max_loras (Optional[int]): Maximum number of LoRA adjustments.
+ vllm_max_cpu_loras (Optional[int]): Maximum number of LoRA adjustments stored on CPU.
+ vllm_lora_extra_vocab_size (int): Additional vocabulary size for LoRA.
+ vllm_placement_group (Optional[dict]): Ray placement group for distributed execution.
+ vllm_log_stats (bool): Whether to log statistics during model operation.
+ notification_email (Optional[str]): Email to send notifications upon completion.
+ batch_size (int): Number of prompts to process in each batch for efficient memory usage.
+ **kwargs: Additional keyword arguments for generation settings like temperature, top_p, etc.
+
+ This method automates the loading of large datasets, generation of text completions, and saving results,
+ facilitating efficient and scalable text generation tasks.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.notification_email = notification_email
+
+ self.model: LLM = self.load_models_vllm(
+ model=model_name,
+ tokenizer=tokenizer_name,
+ tokenizer_mode=vllm_tokenizer_mode,
+ trust_remote_code=True,
+ download_dir=vllm_download_dir,
+ load_format=vllm_load_format,
+ dtype=self._get_torch_dtype(precision),
+ seed=vllm_seed,
+ revision=model_revision,
+ tokenizer_revision=tokenizer_revision,
+ max_model_len=vllm_max_model_len,
+ quantization=(None if quantization == 0 else f"{quantization}-bit"),
+ enforce_eager=vllm_enforce_eager,
+ max_context_len_to_capture=vllm_max_context_len_to_capture,
+ block_size=vllm_block_size,
+ gpu_memory_utilization=vllm_gpu_memory_utilization,
+ swap_space=vllm_swap_space,
+ cache_dtype="auto",
+ sliding_window=vllm_sliding_window,
+ pipeline_parallel_size=vllm_pipeline_parallel_size,
+ tensor_parallel_size=vllm_tensor_parallel_size,
+ worker_use_ray=vllm_worker_use_ray,
+ max_parallel_loading_workers=vllm_max_parallel_loading_workers,
+ disable_custom_all_reduce=vllm_disable_custom_all_reduce,
+ max_num_batched_tokens=vllm_max_num_batched_tokens,
+ max_num_seqs=vllm_max_num_seqs,
+ max_paddings=vllm_max_paddings,
+ device="cuda" if use_cuda else "cpu",
+ max_lora_rank=vllm_max_lora_rank,
+ max_loras=vllm_max_loras,
+ max_cpu_loras=vllm_max_cpu_loras,
+ lora_dtype=self._get_torch_dtype(precision),
+ lora_extra_vocab_size=vllm_lora_extra_vocab_size,
+ placement_group=vllm_placement_group, # type: ignore
+ log_stats=vllm_log_stats,
+ batched_inference=True,
+ )
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["instruction"]
+
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+
+ outputs = self.model.generate(
+ prompts=batch,
+ sampling_params=SamplingParams(
+ n=generation_args.get("n", 1),
+ best_of=generation_args.get("best_of", None),
+ presence_penalty=generation_args.get("presence_penalty", 0.0),
+ frequency_penalty=generation_args.get("frequency_penalty", 0.0),
+ repetition_penalty=generation_args.get("repetition_penalty", 1.0),
+ temperature=generation_args.get("temperature", 1.0),
+ top_p=generation_args.get("top_p", 1.0),
+ top_k=generation_args.get("top_k", -1),
+ min_p=generation_args.get("min_p", 0.0),
+ use_beam_search=generation_args.get("use_beam_search", False),
+ length_penalty=generation_args.get("length_penalty", 1.0),
+ early_stopping=generation_args.get("early_stopping", False),
+ stop=generation_args.get("stop", None),
+ stop_token_ids=generation_args.get("stop_token_ids", None),
+ include_stop_str_in_output=generation_args.get("include_stop_str_in_output", False),
+ ignore_eos=generation_args.get("ignore_eos", False),
+ max_tokens=generation_args.get("max_tokens", 16),
+ logprobs=generation_args.get("logprobs", None),
+ prompt_logprobs=generation_args.get("prompt_logprobs", None),
+ skip_special_tokens=generation_args.get("skip_special_tokens", True),
+ spaces_between_special_tokens=generation_args.get("spaces_between_special_tokens", True),
+ logits_processors=generation_args.get("logits_processors", None),
+ ),
+ )
+ completions = [" ".join(t.text for t in o.outputs) for o in outputs]
+ self._save_completions(completions, batch, output_path)
+ self.done()
+
+ def perform_llama_cpp(
+ self,
+ model: str,
+ filename: Optional[str] = None,
+ local_dir: Optional[Union[str, os.PathLike[str]]] = None,
+ n_gpu_layers: int = 0,
+ split_mode: int = llama_cpp.LLAMA_SPLIT_LAYER,
+ main_gpu: int = 0,
+ tensor_split: Optional[List[float]] = None,
+ vocab_only: bool = False,
+ use_mmap: bool = True,
+ use_mlock: bool = False,
+ kv_overrides: Optional[Dict[str, Union[bool, int, float]]] = None,
+ seed: int = llama_cpp.LLAMA_DEFAULT_SEED,
+ n_ctx: int = 512,
+ n_batch: int = 512,
+ n_threads: Optional[int] = None,
+ n_threads_batch: Optional[int] = None,
+ rope_scaling_type: Optional[int] = llama_cpp.LLAMA_ROPE_SCALING_UNSPECIFIED,
+ rope_freq_base: float = 0.0,
+ rope_freq_scale: float = 0.0,
+ yarn_ext_factor: float = -1.0,
+ yarn_attn_factor: float = 1.0,
+ yarn_beta_fast: float = 32.0,
+ yarn_beta_slow: float = 1.0,
+ yarn_orig_ctx: int = 0,
+ mul_mat_q: bool = True,
+ logits_all: bool = False,
+ embedding: bool = False,
+ offload_kqv: bool = True,
+ last_n_tokens_size: int = 64,
+ lora_base: Optional[str] = None,
+ lora_scale: float = 1.0,
+ lora_path: Optional[str] = None,
+ numa: Union[bool, int] = False,
+ chat_format: Optional[str] = None,
+ chat_handler: Optional[llama_cpp.llama_chat_format.LlamaChatCompletionHandler] = None,
+ draft_model: Optional[llama_cpp.LlamaDraftModel] = None,
+ tokenizer: Optional[PreTrainedTokenizerBase] = None,
+ verbose: bool = True,
+ notification_email: Optional[str] = None,
+ **kwargs,
+ ) -> None:
+ """
+ Performs bulk text generation using the LLaMA model with llama.cpp backend. This method handles the entire
+ process, including model loading, prompt processing, text generation, and saving the results.
+
+ Args:
+ model: Path or identifier for the LLaMA model.
+ filename: Optional filename or glob pattern to match the model file.
+ local_dir: Local directory to save the model files.
+ n_gpu_layers: Number of layers to offload to GPU.
+ split_mode: Split mode for distributing model across GPUs.
+ main_gpu: Main GPU index.
+ tensor_split: Configuration for tensor splitting across GPUs.
+ vocab_only: Whether to load only the vocabulary.
+ use_mmap: Use memory-mapped files for model loading.
+ use_mlock: Lock model data in RAM to prevent swapping.
+ kv_overrides: Key-value pairs for overriding model config.
+ seed: Seed for random number generation.
+ n_ctx: Number of context tokens for generation.
+ n_batch: Batch size for processing.
+ n_threads: Number of threads for generation.
+ n_threads_batch: Number of threads for batch processing.
+ rope_scaling_type: Scaling type for RoPE.
+ rope_freq_base: Base frequency for RoPE.
+ rope_freq_scale: Frequency scaling for RoPE.
+ yarn_ext_factor: YaRN extrapolation factor.
+ yarn_attn_factor: YaRN attention factor.
+ yarn_beta_fast: YaRN beta fast parameter.
+ yarn_beta_slow: YaRN beta slow parameter.
+ yarn_orig_ctx: Original context size for YaRN.
+ mul_mat_q: Multiply matrices for queries.
+ logits_all: Return logits for all tokens.
+ embedding: Enable embedding mode.
+ offload_kqv: Offload K, Q, V matrices to GPU.
+ last_n_tokens_size: Size for the last_n_tokens buffer.
+ lora_base: Base model path for LoRA.
+ lora_scale: Scale factor for LoRA adjustments.
+ lora_path: Path for LoRA adjustments.
+ numa: NUMA configuration.
+ chat_format: Chat format configuration.
+ chat_handler: Handler for chat completions.
+ draft_model: Draft model for speculative decoding.
+ tokenizer: Custom tokenizer instance.
+ verbose: Enable verbose logging.
+ notification_email (Optional[str]): Email to send notifications upon completion.
+ **kwargs: Additional arguments for model loading and text generation.
+ """
+ self.notification_email = notification_email
+
+ # Loading the LLaMA model with llama.cpp
+ llama_model, custom_tokenizer = self.load_models_llama_cpp(
+ model=model,
+ filename=filename,
+ local_dir=local_dir,
+ n_gpu_layers=n_gpu_layers,
+ split_mode=split_mode,
+ main_gpu=main_gpu,
+ tensor_split=tensor_split,
+ vocab_only=vocab_only,
+ use_mmap=use_mmap,
+ use_mlock=use_mlock,
+ kv_overrides=kv_overrides,
+ seed=seed,
+ n_ctx=n_ctx,
+ n_batch=n_batch,
+ n_threads=n_threads,
+ n_threads_batch=n_threads_batch,
+ rope_scaling_type=rope_scaling_type,
+ rope_freq_base=rope_freq_base,
+ rope_freq_scale=rope_freq_scale,
+ yarn_ext_factor=yarn_ext_factor,
+ yarn_attn_factor=yarn_attn_factor,
+ yarn_beta_fast=yarn_beta_fast,
+ yarn_beta_slow=yarn_beta_slow,
+ yarn_orig_ctx=yarn_orig_ctx,
+ mul_mat_q=mul_mat_q,
+ logits_all=logits_all,
+ embedding=embedding,
+ offload_kqv=offload_kqv,
+ last_n_tokens_size=last_n_tokens_size,
+ lora_base=lora_base,
+ lora_scale=lora_scale,
+ lora_path=lora_path,
+ numa=numa,
+ chat_format=chat_format,
+ chat_handler=chat_handler,
+ draft_model=draft_model,
+ tokenizer=tokenizer,
+ verbose=verbose,
+ **kwargs,
+ )
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["instruction"]
+
+ for i in range(0, len(dataset), n_batch):
+ batch = dataset[i : i + n_batch]
+
+ # Generate completion for each prompt using llama_model
+ completions = llama_model.create_chat_completion(
+ messages=[{"content": prompt, "role": "user"} for prompt in batch],
+ functions=generation_args.get("functions"),
+ function_call=generation_args.get("function_call"),
+ tools=generation_args.get("tools"),
+ tool_choice=generation_args.get("tool_choice"),
+ temperature=generation_args.get("temperature", 0.2),
+ top_p=generation_args.get("top_p", 0.95),
+ top_k=generation_args.get("top_k", 40),
+ min_p=generation_args.get("min_p", 0.05),
+ typical_p=generation_args.get("typical_p", 1.0),
+ stream=generation_args.get("stream", False),
+ stop=generation_args.get("stop", []),
+ seed=generation_args.get("seed"),
+ response_format=generation_args.get("response_format"),
+ max_tokens=generation_args.get("max_tokens"),
+ presence_penalty=generation_args.get("presence_penalty", 0.0),
+ frequency_penalty=generation_args.get("frequency_penalty", 0.0),
+ repeat_penalty=generation_args.get("repeat_penalty", 1.1),
+ tfs_z=generation_args.get("tfs_z", 1.0),
+ mirostat_mode=generation_args.get("mirostat_mode", 0),
+ mirostat_tau=generation_args.get("mirostat_tau", 5.0),
+ mirostat_eta=generation_args.get("mirostat_eta", 0.1),
+ model=generation_args.get("model"),
+ logits_processor=generation_args.get("logits_processor"),
+ grammar=generation_args.get("grammar"),
+ logit_bias=generation_args.get("logit_bias"),
+ logprobs=generation_args.get("logprobs"),
+ top_logprobs=generation_args.get("top_logprobs"),
+ )
+
+ self._save_completions(completions["choices"][0]["message"]["content"], batch, output_path) # type: ignore
+ self.done()
+
+ def _save_completions(self, completions: List[str], prompts: List[str], output_path: str) -> None:
+ """
+ Saves the generated texts alongside their prompts to the specified output path. This method ensures the results
+ of text generation are persisted for later use or analysis.
+
+ Args:
+ completions (List[str]): The list of generated texts.
+ prompts (List[str]): The list of prompts corresponding to the generated texts.
+ output_path (str): The directory path to save the results.
+ """
+ data_to_save = [
+ {"prompt": prompt, "completion": completion} for prompt, completion in zip(prompts, completions)
+ ]
+ with open(os.path.join(output_path, f"completions-{str(uuid.uuid4())}.json"), "w") as f:
+ json.dump(data_to_save, f)
diff --git a/huggingface/instruction_tuning.py b/geniusrise_text/instruction/fine_tune.py
similarity index 67%
rename from huggingface/instruction_tuning.py
rename to geniusrise_text/instruction/fine_tune.py
index 8106696..b4bdbcb 100644
--- a/huggingface/instruction_tuning.py
+++ b/geniusrise_text/instruction/fine_tune.py
@@ -1,22 +1,20 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ The text content
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'text' key.
+ ```yaml
+ - text: "The text content"
+ ```
+
+ ### TSV
+ Should contain 'text' column separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'text' column.
+
+ ### SQLite (.db)
+ Should contain a table with 'text' column.
+
+ ### Feather
+ Should contain 'text' column.
+ """
+
+ self.max_length = max_length
+
+ if hasattr(self, "tokenizer") and self.tokenizer is not None:
+ self.label_to_id = self.model.config.label2id if self.model and self.model.config.label2id else {} # type: ignore
+
+ try:
+ self.log.info(f"Loading dataset from {dataset_path}")
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ # Load dataset saved by Hugging Face datasets library
+ return load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ text = record.find("text").text # type: ignore
+ data.append({"text": text})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT text FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ data = [fn(d) for d in data]
+ else:
+ data = data
+
+ return Dataset.from_pandas(pd.DataFrame(data))
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def complete(
+ self,
+ model_name: str,
+ model_class: str = "AutoModelForCausalLM",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ decoding_strategy: str = "generate",
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Performs text completion on the loaded dataset using the specified model and tokenizer. The method handles the
+ entire process, including model loading, text generation, and saving the results.
+
+ Args:
+ model_name (str): The name of the language model to use for text completion.
+ model_class (str, optional): The class of the language model. Defaults to "AutoModelForCausalLM".
+ tokenizer_class (str, optional): The class of the tokenizer. Defaults to "AutoTokenizer".
+ use_cuda (bool, optional): Whether to use CUDA for model inference. Defaults to False.
+ precision (str, optional): Precision for model computation. Defaults to "float16".
+ quantization (int, optional): Level of quantization for optimizing model size and speed. Defaults to 0.
+ device_map (str | Dict | None, optional): Specific device to use for computation. Defaults to "auto".
+ max_memory (Dict, optional): Maximum memory configuration for devices. Defaults to {0: "24GB"}.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool, optional): Whether to enable AWQ optimization. Defaults to False.
+ flash_attention (bool, optional): Whether to use flash attention optimization. Defaults to False.
+ decoding_strategy (str, optional): Strategy for decoding the completion. Defaults to "generate".
+ **kwargs: Additional keyword arguments for text generation.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["text"]
+
+ prompts = []
+ completions = []
+ for _, prompt in enumerate(dataset):
+ completion = self.generate(
+ prompt=prompt,
+ decoding_strategy=decoding_strategy,
+ **generation_args,
+ )
+ completions.append(completion)
+ prompts.append(prompt)
+
+ self._save_completions(completions, prompts, output_path)
+ self.done()
+
+ def complete_vllm(
+ self,
+ model_name: str,
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ # VLLM params
+ vllm_tokenizer_mode: str = "auto",
+ vllm_download_dir: Optional[str] = None,
+ vllm_load_format: str = "auto",
+ vllm_seed: int = 42,
+ vllm_max_model_len: int = 1024,
+ vllm_enforce_eager: bool = False,
+ vllm_max_context_len_to_capture: int = 8192,
+ vllm_block_size: int = 16,
+ vllm_gpu_memory_utilization: float = 0.90,
+ vllm_swap_space: int = 4,
+ vllm_sliding_window: Optional[int] = None,
+ vllm_pipeline_parallel_size: int = 1,
+ vllm_tensor_parallel_size: int = 1,
+ vllm_worker_use_ray: bool = False,
+ vllm_max_parallel_loading_workers: Optional[int] = None,
+ vllm_disable_custom_all_reduce: bool = False,
+ vllm_max_num_batched_tokens: Optional[int] = None,
+ vllm_max_num_seqs: int = 64,
+ vllm_max_paddings: int = 512,
+ vllm_max_lora_rank: Optional[int] = None,
+ vllm_max_loras: Optional[int] = None,
+ vllm_max_cpu_loras: Optional[int] = None,
+ vllm_lora_extra_vocab_size: int = 0,
+ vllm_placement_group: Optional[dict] = None,
+ vllm_log_stats: bool = False,
+ # Generate params
+ notification_email: Optional[str] = None,
+ batch_size: int = 32,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Performs bulk text generation using the Versatile Language Learning Model (VLLM) with specified parameters
+ for fine-tuning model behavior, including quantization and parallel processing settings. This method is designed
+ to process large datasets efficiently by leveraging VLLM capabilities for generating high-quality text completions
+ based on provided prompts.
+
+ Args:
+ model_name (str): The name or path of the VLLM model to use for text generation.
+ use_cuda (bool): Flag indicating whether to use CUDA for GPU acceleration.
+ precision (str): Precision of computations, can be "float16", "bfloat16", etc.
+ quantization (int): Level of quantization for model weights, 0 for none.
+ device_map (str | Dict | None): Specific device(s) to use for model inference.
+ vllm_tokenizer_mode (str): Mode of the tokenizer ("auto", "fast", or "slow").
+ vllm_download_dir (Optional[str]): Directory to download and load the model and tokenizer.
+ vllm_load_format (str): Format to load the model, e.g., "auto", "pt".
+ vllm_seed (int): Seed for random number generation.
+ vllm_max_model_len (int): Maximum sequence length the model can handle.
+ vllm_enforce_eager (bool): Enforce eager execution instead of using optimization techniques.
+ vllm_max_context_len_to_capture (int): Maximum context length for CUDA graph capture.
+ vllm_block_size (int): Block size for caching mechanism.
+ vllm_gpu_memory_utilization (float): Fraction of GPU memory to use.
+ vllm_swap_space (int): Amount of swap space to use in GiB.
+ vllm_sliding_window (Optional[int]): Size of the sliding window for processing.
+ vllm_pipeline_parallel_size (int): Number of pipeline parallel groups.
+ vllm_tensor_parallel_size (int): Number of tensor parallel groups.
+ vllm_worker_use_ray (bool): Whether to use Ray for model workers.
+ vllm_max_parallel_loading_workers (Optional[int]): Maximum number of workers for parallel loading.
+ vllm_disable_custom_all_reduce (bool): Disable custom all-reduce kernel and fall back to NCCL.
+ vllm_max_num_batched_tokens (Optional[int]): Maximum number of tokens to be processed in a single iteration.
+ vllm_max_num_seqs (int): Maximum number of sequences to be processed in a single iteration.
+ vllm_max_paddings (int): Maximum number of paddings to be added to a batch.
+ vllm_max_lora_rank (Optional[int]): Maximum rank for LoRA adjustments.
+ vllm_max_loras (Optional[int]): Maximum number of LoRA adjustments.
+ vllm_max_cpu_loras (Optional[int]): Maximum number of LoRA adjustments stored on CPU.
+ vllm_lora_extra_vocab_size (int): Additional vocabulary size for LoRA.
+ vllm_placement_group (Optional[dict]): Ray placement group for distributed execution.
+ vllm_log_stats (bool): Whether to log statistics during model operation.
+ notification_email (Optional[str]): Email to send notifications upon completion.
+ batch_size (int): Number of prompts to process in each batch for efficient memory usage.
+ **kwargs: Additional keyword arguments for generation settings like temperature, top_p, etc.
+
+ This method automates the loading of large datasets, generation of text completions, and saving results,
+ facilitating efficient and scalable text generation tasks.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.notification_email = notification_email
+
+ self.model: LLM = self.load_models_vllm(
+ model=model_name,
+ tokenizer=tokenizer_name,
+ tokenizer_mode=vllm_tokenizer_mode,
+ trust_remote_code=True,
+ download_dir=vllm_download_dir,
+ load_format=vllm_load_format,
+ dtype=self._get_torch_dtype(precision),
+ seed=vllm_seed,
+ revision=model_revision,
+ tokenizer_revision=tokenizer_revision,
+ max_model_len=vllm_max_model_len,
+ quantization=(None if quantization == 0 else f"{quantization}-bit"),
+ enforce_eager=vllm_enforce_eager,
+ max_context_len_to_capture=vllm_max_context_len_to_capture,
+ block_size=vllm_block_size,
+ gpu_memory_utilization=vllm_gpu_memory_utilization,
+ swap_space=vllm_swap_space,
+ cache_dtype="auto",
+ sliding_window=vllm_sliding_window,
+ pipeline_parallel_size=vllm_pipeline_parallel_size,
+ tensor_parallel_size=vllm_tensor_parallel_size,
+ worker_use_ray=vllm_worker_use_ray,
+ max_parallel_loading_workers=vllm_max_parallel_loading_workers,
+ disable_custom_all_reduce=vllm_disable_custom_all_reduce,
+ max_num_batched_tokens=vllm_max_num_batched_tokens,
+ max_num_seqs=vllm_max_num_seqs,
+ max_paddings=vllm_max_paddings,
+ device="cuda" if use_cuda else "cpu",
+ max_lora_rank=vllm_max_lora_rank,
+ max_loras=vllm_max_loras,
+ max_cpu_loras=vllm_max_cpu_loras,
+ lora_dtype=self._get_torch_dtype(precision),
+ lora_extra_vocab_size=vllm_lora_extra_vocab_size,
+ placement_group=vllm_placement_group, # type: ignore
+ log_stats=vllm_log_stats,
+ batched_inference=True,
+ )
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["text"]
+
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+
+ outputs = self.model.generate(
+ prompts=batch,
+ sampling_params=SamplingParams(
+ n=generation_args.get("n", 1),
+ best_of=generation_args.get("best_of", None),
+ presence_penalty=generation_args.get("presence_penalty", 0.0),
+ frequency_penalty=generation_args.get("frequency_penalty", 0.0),
+ repetition_penalty=generation_args.get("repetition_penalty", 1.0),
+ temperature=generation_args.get("temperature", 1.0),
+ top_p=generation_args.get("top_p", 1.0),
+ top_k=generation_args.get("top_k", -1),
+ min_p=generation_args.get("min_p", 0.0),
+ use_beam_search=generation_args.get("use_beam_search", False),
+ length_penalty=generation_args.get("length_penalty", 1.0),
+ early_stopping=generation_args.get("early_stopping", False),
+ stop=generation_args.get("stop", None),
+ stop_token_ids=generation_args.get("stop_token_ids", None),
+ include_stop_str_in_output=generation_args.get("include_stop_str_in_output", False),
+ ignore_eos=generation_args.get("ignore_eos", False),
+ max_tokens=generation_args.get("max_tokens", 16),
+ logprobs=generation_args.get("logprobs", None),
+ prompt_logprobs=generation_args.get("prompt_logprobs", None),
+ skip_special_tokens=generation_args.get("skip_special_tokens", True),
+ spaces_between_special_tokens=generation_args.get("spaces_between_special_tokens", True),
+ logits_processors=generation_args.get("logits_processors", None),
+ ),
+ )
+ completions = [" ".join(t.text for t in o.outputs) for o in outputs]
+ self._save_completions(completions, batch, output_path)
+ self.done()
+
+ def complete_llama_cpp(
+ self,
+ model: str,
+ filename: Optional[str] = None,
+ local_dir: Optional[Union[str, os.PathLike[str]]] = None,
+ n_gpu_layers: int = 0,
+ split_mode: int = llama_cpp.LLAMA_SPLIT_LAYER,
+ main_gpu: int = 0,
+ tensor_split: Optional[List[float]] = None,
+ vocab_only: bool = False,
+ use_mmap: bool = True,
+ use_mlock: bool = False,
+ kv_overrides: Optional[Dict[str, Union[bool, int, float]]] = None,
+ seed: int = llama_cpp.LLAMA_DEFAULT_SEED,
+ n_ctx: int = 512,
+ n_batch: int = 512,
+ n_threads: Optional[int] = None,
+ n_threads_batch: Optional[int] = None,
+ rope_scaling_type: Optional[int] = llama_cpp.LLAMA_ROPE_SCALING_UNSPECIFIED,
+ rope_freq_base: float = 0.0,
+ rope_freq_scale: float = 0.0,
+ yarn_ext_factor: float = -1.0,
+ yarn_attn_factor: float = 1.0,
+ yarn_beta_fast: float = 32.0,
+ yarn_beta_slow: float = 1.0,
+ yarn_orig_ctx: int = 0,
+ mul_mat_q: bool = True,
+ logits_all: bool = False,
+ embedding: bool = False,
+ offload_kqv: bool = True,
+ last_n_tokens_size: int = 64,
+ lora_base: Optional[str] = None,
+ lora_scale: float = 1.0,
+ lora_path: Optional[str] = None,
+ numa: Union[bool, int] = False,
+ chat_format: Optional[str] = None,
+ chat_handler: Optional[llama_cpp.llama_chat_format.LlamaChatCompletionHandler] = None,
+ draft_model: Optional[llama_cpp.LlamaDraftModel] = None,
+ tokenizer: Optional[PreTrainedTokenizerBase] = None,
+ verbose: bool = True,
+ notification_email: Optional[str] = None,
+ **kwargs,
+ ) -> None:
+ """
+ Performs bulk text generation using the LLaMA model with llama.cpp backend. This method handles the entire
+ process, including model loading, prompt processing, text generation, and saving the results.
+
+ Args:
+ model: Path or identifier for the LLaMA model.
+ filename: Optional filename or glob pattern to match the model file.
+ local_dir: Local directory to save the model files.
+ n_gpu_layers: Number of layers to offload to GPU.
+ split_mode: Split mode for distributing model across GPUs.
+ main_gpu: Main GPU index.
+ tensor_split: Configuration for tensor splitting across GPUs.
+ vocab_only: Whether to load only the vocabulary.
+ use_mmap: Use memory-mapped files for model loading.
+ use_mlock: Lock model data in RAM to prevent swapping.
+ kv_overrides: Key-value pairs for overriding model config.
+ seed: Seed for random number generation.
+ n_ctx: Number of context tokens for generation.
+ n_batch: Batch size for processing.
+ n_threads: Number of threads for generation.
+ n_threads_batch: Number of threads for batch processing.
+ rope_scaling_type: Scaling type for RoPE.
+ rope_freq_base: Base frequency for RoPE.
+ rope_freq_scale: Frequency scaling for RoPE.
+ yarn_ext_factor: YaRN extrapolation factor.
+ yarn_attn_factor: YaRN attention factor.
+ yarn_beta_fast: YaRN beta fast parameter.
+ yarn_beta_slow: YaRN beta slow parameter.
+ yarn_orig_ctx: Original context size for YaRN.
+ mul_mat_q: Multiply matrices for queries.
+ logits_all: Return logits for all tokens.
+ embedding: Enable embedding mode.
+ offload_kqv: Offload K, Q, V matrices to GPU.
+ last_n_tokens_size: Size for the last_n_tokens buffer.
+ lora_base: Base model path for LoRA.
+ lora_scale: Scale factor for LoRA adjustments.
+ lora_path: Path for LoRA adjustments.
+ numa: NUMA configuration.
+ chat_format: Chat format configuration.
+ chat_handler: Handler for chat completions.
+ draft_model: Draft model for speculative decoding.
+ tokenizer: Custom tokenizer instance.
+ verbose: Enable verbose logging.
+ notification_email (Optional[str]): Email to send notifications upon completion.
+ **kwargs: Additional arguments for model loading and text generation.
+ """
+ self.notification_email = notification_email
+
+ # Loading the LLaMA model with llama.cpp
+ llama_model, custom_tokenizer = self.load_models_llama_cpp(
+ model=model,
+ filename=filename,
+ local_dir=local_dir,
+ n_gpu_layers=n_gpu_layers,
+ split_mode=split_mode,
+ main_gpu=main_gpu,
+ tensor_split=tensor_split,
+ vocab_only=vocab_only,
+ use_mmap=use_mmap,
+ use_mlock=use_mlock,
+ kv_overrides=kv_overrides,
+ seed=seed,
+ n_ctx=n_ctx,
+ n_batch=n_batch,
+ n_threads=n_threads,
+ n_threads_batch=n_threads_batch,
+ rope_scaling_type=rope_scaling_type,
+ rope_freq_base=rope_freq_base,
+ rope_freq_scale=rope_freq_scale,
+ yarn_ext_factor=yarn_ext_factor,
+ yarn_attn_factor=yarn_attn_factor,
+ yarn_beta_fast=yarn_beta_fast,
+ yarn_beta_slow=yarn_beta_slow,
+ yarn_orig_ctx=yarn_orig_ctx,
+ mul_mat_q=mul_mat_q,
+ logits_all=logits_all,
+ embedding=embedding,
+ offload_kqv=offload_kqv,
+ last_n_tokens_size=last_n_tokens_size,
+ lora_base=lora_base,
+ lora_scale=lora_scale,
+ lora_path=lora_path,
+ numa=numa,
+ chat_format=chat_format,
+ chat_handler=chat_handler,
+ draft_model=draft_model,
+ tokenizer=tokenizer,
+ verbose=verbose,
+ **kwargs,
+ )
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["instruction"]
+
+ for i in range(0, len(dataset), n_batch):
+ batch = dataset[i : i + n_batch]
+ completions = []
+
+ for prompt in batch:
+ # Generate completion for each prompt using llama_model
+ completion = llama_model.create_completion(
+ prompt=prompt,
+ suffix=generation_args.get("suffix", None),
+ max_tokens=generation_args.get("max_tokens", 16),
+ temperature=generation_args.get("temperature", 0.8),
+ top_p=generation_args.get("top_p", 0.95),
+ min_p=generation_args.get("min_p", 0.05),
+ typical_p=generation_args.get("typical_p", 1.0),
+ logprobs=generation_args.get("logprobs", None),
+ echo=generation_args.get("echo", False),
+ stop=generation_args.get("stop", []),
+ frequency_penalty=generation_args.get("frequency_penalty", 0.0),
+ presence_penalty=generation_args.get("presence_penalty", 0.0),
+ repeat_penalty=generation_args.get("repeat_penalty", 1.1),
+ top_k=generation_args.get("top_k", 40),
+ seed=generation_args.get("seed", None),
+ tfs_z=generation_args.get("tfs_z", 1.0),
+ mirostat_mode=generation_args.get("mirostat_mode", 0),
+ mirostat_tau=generation_args.get("mirostat_tau", 5.0),
+ mirostat_eta=generation_args.get("mirostat_eta", 0.1),
+ model=generation_args.get("model", None),
+ stopping_criteria=generation_args.get("stopping_criteria", None),
+ logits_processor=generation_args.get("logits_processor", None),
+ grammar=generation_args.get("grammar", None),
+ logit_bias=generation_args.get("logit_bias", None),
+ )
+ completions.append(completion)
+
+ self._save_completions([c["choices"][0]["text"] for c in completions], batch, output_path) # type: ignore
+ self.done()
+
+ def _save_completions(self, completions: List[str], prompts: List[str], output_path: str) -> None:
+ """
+ Saves the generated completions to the specified output path.
+
+ Args:
+ completions (List[str]): The list of generated text completions.
+ prompts (List[str]): The list of prompts corresponding to the completions.
+ output_path (str): The path to save the completion results.
+
+ This method is called internally by the complete method to persist the completion results.
+ """
+ data_to_save = [
+ {"prompt": prompt, "completion": completion} for prompt, completion in zip(prompts, completions)
+ ]
+ with open(os.path.join(output_path, f"completions-{str(uuid.uuid4())}.json"), "w") as f:
+ json.dump(data_to_save, f)
diff --git a/huggingface/language_model.py b/geniusrise_text/language_model/fine_tune.py
similarity index 73%
rename from huggingface/language_model.py
rename to geniusrise_text/language_model/fine_tune.py
index 105f7ff..c81862b 100644
--- a/huggingface/language_model.py
+++ b/geniusrise_text/language_model/fine_tune.py
@@ -1,19 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ token1 token2 ...
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'tokens' keys.
+ ```yaml
+ - tokens: ["token1", "token2", ...]
+ ```
+
+ ### TSV
+ Should contain 'tokens' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'tokens' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'tokens' columns.
+
+ ### Feather
+ Should contain 'tokens' columns.
+ """
+ self.log.info(f"Loading dataset from {dataset_path}")
+ try:
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ return load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ text = record.find("text").text.split() # type: ignore
+ data.append({"text": text})
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT text FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ return Dataset.from_pandas(pd.DataFrame(data))
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def recognize_entities(
+ self,
+ model_name: str,
+ max_length: int = 512,
+ model_class: str = "AutoModelForSeq2SeqLM",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Performs bulk named entity recognition on the loaded dataset. The method processes the text in batches,
+ applying the NER model to recognize entities.
+
+ Args:
+ model_name (str): The name or path of the NER model.
+ max_length (int): The maximum sequence length for the tokenizer.
+ model_class (str): The class of the model, defaults to "AutoModelForTokenClassification".
+ tokenizer_class (str): The class of the tokenizer, defaults to "AutoTokenizer".
+ use_cuda (bool): Whether to use CUDA for model inference, defaults to False.
+ precision (str): Model computation precision, defaults to "float16".
+ quantization (int): Level of quantization for model size and speed optimization, defaults to 0.
+ device_map (str | Dict | None): Specific device configuration for computation, defaults to "auto".
+ max_memory (Dict): Maximum memory configuration for the devices.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool): Whether to enable AWQ optimization, defaults to False.
+ flash_attention (bool): Whether to use flash attention optimization, defaults to False.
+ batch_size (int): Number of documents to process simultaneously, defaults to 32.
+ **kwargs: Arbitrary keyword arguments for additional configuration.
+
+ Returns:
+ None: The method processes the dataset and saves the predictions without returning any value.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ dataset = self.load_dataset(dataset_path)
+ if dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ if dataset:
+ dataset = dataset["text"]
+
+ # Process data in batches
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+ inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
+
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+
+ predictions = self.model(**inputs, **generation_args)
+ predictions = predictions[0] if isinstance(predictions, tuple) else predictions.logits
+ predictions = predictions.argmax(dim=-1).squeeze().tolist()
+
+ self._save_predictions(inputs["input_ids"].tolist(), predictions, batch, output_path, i)
+ self.done()
+
+ def _save_predictions(
+ self, inputs: list, predictions: list, input_batch: List[str], output_path: str, batch_idx: int
+ ) -> None:
+ """
+ Saves the NER predictions to the specified output path.
+
+ Args:
+ inputs (list): List of input tokens.
+ predictions (list): List of prediction tensors from the NER model.
+ input_batch (List[str]): The input text batch.
+ output_path (str): The path to save the prediction results.
+ batch_idx (int): The index of the current batch, used for naming the output files.
+
+ Returns:
+ None: The method saves the predictions to files and does not return any value.
+ """
+ # Convert tensor of label ids to list of label strings
+ label_predictions = [
+ [
+ {
+ "label": self.model.config.id2label[label_id],
+ "position": i,
+ "token": self.tokenizer.convert_ids_to_tokens(inp[i]),
+ }
+ for i, label_id in enumerate(pred)
+ ]
+ for pred, inp in zip(predictions, inputs)
+ ]
+
+ # Prepare data for saving
+ data_to_save = [
+ {"input": input_text, "labels": label} for input_text, label in zip(input_batch, label_predictions)
+ ]
+ with open(os.path.join(output_path, f"predictions-{batch_idx}-{str(uuid.uuid4())}.jsonl"), "w") as f:
+ for item in data_to_save:
+ f.write(json.dumps(item) + "\n")
+
+ self.log.info(f"Saved predictions for batch {batch_idx} to {output_path}")
diff --git a/huggingface/ner.py b/geniusrise_text/ner/fine_tune.py
similarity index 67%
rename from huggingface/ner.py
rename to geniusrise_text/ner/fine_tune.py
index 95da051..1f59b7f 100644
--- a/huggingface/ner.py
+++ b/geniusrise_text/ner/fine_tune.py
@@ -1,18 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ The premise text
+ The hypothesis text
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'premise' and 'hypothesis' keys.
+ ```yaml
+ - premise: "The premise text"
+ hypothesis: "The hypothesis text"
+ ```
+
+ ### TSV
+ Should contain 'premise' and 'hypothesis' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'premise' and 'hypothesis' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'premise' and 'hypothesis' columns.
+
+ ### Feather
+ Should contain 'premise' and 'hypothesis' columns.
+ """
+ self.max_length = max_length
+
+ self.log.info(f"Loading dataset from {dataset_path}")
+ try:
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ dataset = load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ premise = record.find("premise").text # type: ignore
+ hypothesis = record.find("hypothesis").text # type: ignore
+ data.append({"premise": premise, "hypothesis": hypothesis})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT premise, hypothesis FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ return Dataset.from_pandas(pd.DataFrame(data))
+
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def infer(
+ self,
+ model_name: str,
+ max_length: int = 512,
+ model_class: str = "AutoModelForSeq2SeqLM",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Performs NLI inference on a loaded dataset using the specified model. The method processes the data in batches and saves
+ the results to the configured output path.
+
+ Args:
+ model_name (str): Name or path of the NLI model.
+ max_length (int, optional): Maximum length of the sequences for tokenization purposes. Defaults to 512.
+ model_class (str, optional): Class name of the model (e.g., "AutoModelForSequenceClassification"). Defaults to "AutoModelForSeq2SeqLM".
+ tokenizer_class (str, optional): Class name of the tokenizer (e.g., "AutoTokenizer"). Defaults to "AutoTokenizer".
+ use_cuda (bool, optional): Whether to use CUDA for model inference. Defaults to False.
+ precision (str, optional): Precision for model computation (e.g., "float16"). Defaults to "float16".
+ quantization (int, optional): Level of quantization for optimizing model size and speed. Defaults to 0.
+ device_map (str | Dict | None, optional): Specific device to use for computation. Defaults to "auto".
+ max_memory (Dict, optional): Maximum memory configuration for devices. Defaults to {0: "24GB"}.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool, optional): Whether to enable AWQ optimization. Defaults to False.
+ flash_attention (bool, optional): Whether to use flash attention optimization. Defaults to False.
+ batch_size (int, optional): Number of premise-hypothesis pairs to process simultaneously. Defaults to 32.
+ **kwargs: Arbitrary keyword arguments for model and generation configurations.
+ ```
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ dataset = self.load_dataset(dataset_path)
+ if dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+
+ predictions = []
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+
+ inputs = self.tokenizer(
+ batch["premise"],
+ batch["hypothesis"],
+ padding=True,
+ return_tensors="pt",
+ )
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+
+ with torch.no_grad():
+ outputs = self.model(**inputs)
+ logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
+ if next(self.model.parameters()).is_cuda:
+ logits = logits.cpu()
+ softmax = torch.nn.functional.softmax(logits, dim=-1)
+ scores = softmax.numpy().tolist()
+
+ for score in scores:
+ label_scores = {
+ self.model.config.id2label[label_id]: score for label_id, score in enumerate(scores[0])
+ }
+ predictions.append(label_scores)
+
+ # Save results
+ self.log.info(f"Saving results to {output_path}")
+ os.makedirs(output_path, exist_ok=True)
+ output_file = os.path.join(output_path, f"nli_results_{uuid.uuid4().hex}.jsonl")
+ with open(output_file, "w") as f:
+ for data, pred in zip(dataset, predictions):
+ result = {
+ "premise": data["premise"],
+ "hypothesis": data["hypothesis"],
+ "prediction": pred,
+ }
+ f.write(json.dumps(result) + "\n")
+
+ self.done()
+ self.log.info("Inference completed.")
diff --git a/huggingface/commonsense_reasoning.py b/geniusrise_text/nli/fine_tune.py
similarity index 64%
rename from huggingface/commonsense_reasoning.py
rename to geniusrise_text/nli/fine_tune.py
index bb48f3b..9695838 100644
--- a/huggingface/commonsense_reasoning.py
+++ b/geniusrise_text/nli/fine_tune.py
@@ -1,18 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ The context content
+ The question
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'context' and 'question' keys.
+ ```yaml
+ - context: "The context content"
+ question: "The question"
+ ```
+
+ ### TSV
+ Should contain 'context' and 'question' columns separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'context' and 'question' columns.
+
+ ### SQLite (.db)
+ Should contain a table with 'context' and 'question' columns.
+
+ ### Feather
+ Should contain 'context' and 'question' columns.
+
+ Args:
+ dataset_path (str): The path to the dataset directory.
+ pad_on_right (bool): Whether to pad on the right.
+ max_length (int): The maximum length of the sequences.
+ doc_stride (int): The document stride.
+ evaluate_squadv2 (bool): Whether to evaluate using SQuAD v2 metrics.
+
+ Returns:
+ Dataset: The loaded dataset.
+ """
+
+ self.max_length = max_length
+
+ try:
+ self.log.info(f"Loading dataset from {dataset_path}")
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ return load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ context = record.find("data").text # type: ignore
+ question = record.find("question").text # type: ignore
+ data.append({"data": context, "question": question})
+
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT data, question FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ return Dataset.from_pandas(pd.DataFrame(data))
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def answer_questions(
+ self,
+ model_name: str,
+ model_class: str = "AutoModelForQuestionAnswering",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Perform bulk question-answering using the specified model and tokenizer. This method can handle various types
+ of QA models including traditional, TAPAS, and TAPEX.
+
+ Args:
+ model_name (str): Name or path of the question-answering model.
+ model_class (str, optional): Class name of the model (e.g., "AutoModelForQuestionAnswering").
+ tokenizer_class (str, optional): Class name of the tokenizer (e.g., "AutoTokenizer").
+ use_cuda (bool, optional): Whether to use CUDA for model inference. Defaults to False.
+ precision (str, optional): Precision for model computation. Defaults to "float16".
+ quantization (int, optional): Level of quantization for optimizing model size and speed. Defaults to 0.
+ device_map (str | Dict | None, optional): Specific device to use for computation. Defaults to "auto".
+ max_memory (Dict, optional): Maximum memory configuration for devices. Defaults to {0: "24GB"}.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool, optional): Whether to enable AWQ optimization. Defaults to False.
+ flash_attention (bool, optional): Whether to use flash attention optimization. Defaults to False.
+ batch_size (int, optional): Number of questions to process simultaneously. Defaults to 32.
+ **kwargs: Arbitrary keyword arguments for model and generation configurations.
+
+ Processing:
+ The method processes the data in batches, utilizing the appropriate model based on the model name
+ and generating answers for the questions provided in the dataset.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ dataset = self.load_dataset(dataset_path)
+ if dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+
+ model_type = "traditional"
+ if "tapas" in self.model_name.lower():
+ model_type = "tapas"
+ elif "tapex" in self.model_name.lower():
+ model_type = "tapex"
+
+ output_data = []
+ for batch in range(0, len(dataset), self.batch_size):
+ batch_data = dataset[batch : batch + self.batch_size]
+
+ if model_type == "traditional":
+ questions = batch_data["question"]
+ contexts = batch_data["data"]
+
+ inputs = self.tokenizer(
+ questions,
+ contexts,
+ add_special_tokens=True,
+ return_tensors="pt",
+ truncation="only_second",
+ max_length=self.max_length,
+ )
+
+ # Move inputs to GPU if CUDA is available
+ if self.use_cuda and torch.cuda.is_available():
+ inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+ outputs = self.model(**inputs)
+
+ answer_start_scores, answer_end_scores = outputs.start_logits, outputs.end_logits
+ answer_start = torch.argmax(answer_start_scores, dim=1)
+ answer_end = torch.argmax(answer_end_scores, dim=1) + 1
+
+ for i in range(outputs.start_logits.shape[0]):
+ answer = self.tokenizer.convert_tokens_to_string(
+ self.tokenizer.convert_ids_to_tokens(
+ inputs["input_ids"][i][int(answer_start[i]) : int(answer_end[i])]
+ )
+ )
+ output_data.append(
+ {
+ "data": contexts[i],
+ "question": questions[i],
+ "answer": answer,
+ }
+ )
+ elif model_type == "tapas":
+ questions = batch_data["question"]
+ tables = [pd.DataFrame.from_dict(json.loads(x)) for x in batch_data["data"]]
+
+ for table, question in zip(tables, questions):
+ inputs = self.tokenizer(table=table, queries=[question], padding="max_length", return_tensors="pt")
+
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+ outputs = self.model(**inputs)
+
+ # Decode the predicted tokens
+ if hasattr(outputs, "logits_aggregation") and outputs.logits_aggregation is not None:
+ (
+ predicted_answer_coordinates,
+ predicted_aggregation_indices,
+ ) = self.tokenizer.convert_logits_to_predictions(
+ {k: v.cpu() for k, v in inputs.items()},
+ outputs.logits.detach().cpu(),
+ outputs.logits_aggregation.detach().cpu(),
+ )
+ else:
+ predicted_answer_coordinates = self.tokenizer.convert_logits_to_predictions(
+ {k: v.cpu() for k, v in inputs.items()},
+ outputs.logits.detach().cpu(),
+ )
+ predicted_aggregation_indices = None
+
+ cell_answers = [
+ self._convert_coordinates_to_answer(table, x) for x in predicted_answer_coordinates[0]
+ ]
+ if type(cell_answers[0]) is list:
+ cell_answers = [y for x in cell_answers for y in x] # type: ignore
+
+ if predicted_aggregation_indices:
+ aggregation_answer = self._convert_aggregation_to_answer(predicted_aggregation_indices[0])
+ else:
+ aggregation_answer = "NONE"
+ output_data.append(
+ {
+ "data": table.to_dict("records"),
+ "question": question,
+ "answers": cell_answers,
+ "aggregation": aggregation_answer,
+ }
+ )
+
+ elif model_type == "tapex":
+ questions = batch_data["question"]
+ tables = [pd.DataFrame.from_dict(json.loads(x)) for x in batch_data["data"]]
+
+ for table, question in zip(tables, questions):
+ encoding = self.tokenizer(table, question, return_tensors="pt")
+ if next(self.model.parameters()).is_cuda:
+ encoding = {k: v.cuda() for k, v in encoding.items()}
+
+ outputs = self.model.generate(**encoding)
+ answers = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+ output_data.append(
+ {
+ "data": table.to_dict("records"),
+ "question": question,
+ "answers": answers,
+ "aggregation": "NONE",
+ }
+ )
+
+ # Save the results
+ output_file = os.path.join(output_path, f"qa_results-{str(uuid.uuid4())}.json")
+ with open(output_file, "w") as file:
+ json.dump(output_data, file)
+ self.done()
+ self.log.info(f"Results saved to {output_file}")
+
+ def _convert_aggregation_to_answer(self, aggregation_index: int) -> str:
+ """
+ Converts the aggregation index predicted by TAPAS into an aggregation operation.
+
+ Args:
+ aggregation_index (int): The index of the aggregation operation.
+
+ Returns:
+ str: The string representation of the aggregation operation.
+ """
+ aggregation_operations = {
+ 0: "NONE",
+ 1: "SUM",
+ 2: "AVERAGE",
+ 3: "COUNT",
+ 4: "MIN",
+ 5: "MAX",
+ 6: "OR",
+ 7: "AND",
+ 8: "CONCAT",
+ 9: "FIRST",
+ 10: "LAST",
+ }
+ return aggregation_operations.get(aggregation_index, "NONE")
+
+ def _convert_coordinates_to_answer(self, table: pd.DataFrame, coordinates: Any) -> List[str]:
+ """
+ Converts the coordinates predicted by TAPAS into an answer string.
+
+ Args:
+ table (pd.DataFrame): The table used for the QA.
+ coordinates (Any): The coordinates of the cells predicted as part of the answer.
+
+ Returns:
+ List[str]: The answer strings.
+ """
+ if type(coordinates) is tuple:
+ coordinates = [coordinates]
+ return [table.iat[coord] for coord in coordinates]
diff --git a/huggingface/question_answering.py b/geniusrise_text/qa/fine_tune.py
similarity index 82%
rename from huggingface/question_answering.py
rename to geniusrise_text/qa/fine_tune.py
index 53e0060..8694a57 100644
--- a/huggingface/question_answering.py
+++ b/geniusrise_text/qa/fine_tune.py
@@ -1,18 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ The text content
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with a 'text' key.
+ ```yaml
+ - text: "The text content"
+ ```
+
+ ### TSV
+ Should contain a 'text' column.
+
+ ### Excel (.xls, .xlsx)
+ Should contain a 'text' column.
+
+ ### SQLite (.db)
+ Should contain a table with a 'text' column.
+
+ ### Feather
+ Should contain a 'text' column.
+ """
+ try:
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ dataset = load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ text = record.find("text").text # type: ignore
+ data.append({"text": text})
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = "SELECT text FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ data = [fn(d) for d in data]
+ else:
+ data = data
+
+ dataset = Dataset.from_pandas(pd.DataFrame(data))
+ return dataset
+
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def summarize(
+ self,
+ model_name: str,
+ model_class: str = "AutoModelForSeq2SeqLM",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ max_length: int = 512,
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Perform bulk summarization using the specified model and tokenizer. This method handles the entire summarization
+ process including loading the model, processing input data, generating summarization, and saving the results.
+
+ Args:
+ model_name (str): Name or path of the translation model.
+ origin (str): Source language ISO code.
+ target (str): Target language ISO code.
+ max_length (int): Maximum length of the tokens (default 512).
+ model_class (str): Class name of the model (default "AutoModelForSeq2SeqLM").
+ tokenizer_class (str): Class name of the tokenizer (default "AutoTokenizer").
+ use_cuda (bool): Whether to use CUDA for model inference (default False).
+ precision (str): Precision for model computation (default "float16").
+ quantization (int): Level of quantization for optimizing model size and speed (default 0).
+ device_map (str | Dict | None): Specific device to use for computation (default "auto").
+ max_memory (Dict): Maximum memory configuration for devices.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool): Whether to enable AWQ optimization (default False).
+ flash_attention (bool): Whether to use flash attention optimization (default False).
+ batch_size (int): Number of translations to process simultaneously (default 32).
+ max_length (int): Maximum lenght of the summary to be generated (default 512).
+ **kwargs: Arbitrary keyword arguments for model and generation configurations.
+ """
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path, max_lengt=max_length)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset["text"]
+
+ # Process data in batches
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+ inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
+
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+
+ # Generate summaries
+ summaries = self.model.generate(**inputs, **self.generation_args)
+ if next(self.model.parameters()).is_cuda:
+ summaries = summaries.cpu()
+
+ decoded_summaries = [self.tokenizer.decode(s, skip_special_tokens=True) for s in summaries]
+
+ self._save_summaries(decoded_summaries, batch, output_path, i)
+ self.done()
+
+ def _save_summaries(self, summaries: List[str], input_batch: List[str], output_path: str, batch_idx: int) -> None:
+ """
+ Saves the generated summaries to a specified output path. This method is called internally by the summarize method
+ to persist the summarization results.
+
+ Args:
+ summaries (List[str]): List of generated summaries.
+ input_batch (List[str]): List of original texts that were summarized.
+ output_path (str): Path to save the summaries.
+ batch_idx (int): Index of the current batch (for naming files).
+ """
+ data_to_save = [
+ {"input": input_text, "summary": summary} for input_text, summary in zip(input_batch, summaries)
+ ]
+ with open(os.path.join(output_path, f"summaries-{batch_idx}-{str(uuid.uuid4())}.json"), "w") as f:
+ json.dump(data_to_save, f)
diff --git a/huggingface/summarization.py b/geniusrise_text/summarization/fine_tune.py
similarity index 73%
rename from huggingface/summarization.py
rename to geniusrise_text/summarization/fine_tune.py
index d20b39b..869e567 100644
--- a/huggingface/summarization.py
+++ b/geniusrise_text/summarization/fine_tune.py
@@ -1,18 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
+ English text
+
+ ```
+
+ ### YAML
+ Each document should be a dictionary with 'en' key.
+ ```yaml
+ - en: "English text"
+ ```
+
+ ### TSV
+ Should contain 'en' column separated by tabs.
+
+ ### Excel (.xls, .xlsx)
+ Should contain 'en' column.
+
+ ### SQLite (.db)
+ Should contain a table with 'en' column.
+
+ ### Feather
+ Should contain 'en' column.
+
+ Args:
+ dataset_path (str): The path to the directory containing the dataset files.
+ max_length (int, optional): The maximum length for tokenization. Defaults to 512.
+ origin (str, optional): The origin language. Defaults to 'en'.
+ target (str, optional): The target language. Defaults to 'hi'.
+ **kwargs: Additional keyword arguments.
+
+ Returns:
+ DatasetDict: The loaded dataset.
+ """
+
+ self.max_length = max_length
+ self.origin = origin
+ self.target = target
+ if self.origin:
+ self.tokenizer.src_lang = self.origin
+
+ try:
+ if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
+ dataset = load_from_disk(dataset_path)
+ else:
+ data = []
+ for filename in glob.glob(f"{dataset_path}/**/*", recursive=True):
+ filepath = os.path.join(dataset_path, filename)
+ if filename.endswith(".jsonl"):
+ with open(filepath, "r") as f:
+ for line in f:
+ example = json.loads(line)
+ data.append(example)
+ elif filename.endswith(".csv"):
+ df = pd.read_csv(filepath)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".parquet"):
+ df = pq.read_table(filepath).to_pandas()
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".json"):
+ with open(filepath, "r") as f:
+ json_data = json.load(f)
+ data.extend(json_data)
+ elif filename.endswith(".xml"):
+ tree = ET.parse(filepath)
+ root = tree.getroot()
+ for record in root.findall("record"):
+ origin = record.find(self.origin).text # type: ignore
+ data.append({"translation": {self.origin: origin}})
+ elif filename.endswith(".yaml") or filename.endswith(".yml"):
+ with open(filepath, "r") as f:
+ yaml_data = yaml.safe_load(f)
+ data.extend(yaml_data)
+ elif filename.endswith(".tsv"):
+ df = pd.read_csv(filepath, sep="\t")
+ data.extend(df.to_dict("records"))
+ elif filename.endswith((".xls", ".xlsx")):
+ df = pd.read_excel(filepath)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".db"):
+ conn = sqlite3.connect(filepath)
+ query = f"SELECT {self.origin} FROM dataset_table;"
+ df = pd.read_sql_query(query, conn)
+ data.extend(df.to_dict("records"))
+ elif filename.endswith(".feather"):
+ df = feather.read_feather(filepath)
+ data.extend(df.to_dict("records"))
+
+ if hasattr(self, "map_data") and self.map_data:
+ fn = eval(self.map_data) # type: ignore
+ data = [fn(d) for d in data]
+ else:
+ data = data
+
+ dataset = Dataset.from_pandas(pd.DataFrame(data))
+
+ return dataset
+ except Exception as e:
+ self.log.exception(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
+ raise
+
+ def translate(
+ self,
+ model_name: str,
+ origin: str,
+ target: str,
+ max_length: int = 512,
+ model_class: str = "AutoModelForSeq2SeqLM",
+ tokenizer_class: str = "AutoTokenizer",
+ use_cuda: bool = False,
+ precision: str = "float16",
+ quantization: int = 0,
+ device_map: str | Dict | None = "auto",
+ max_memory={0: "24GB"},
+ torchscript: bool = False,
+ compile: bool = False,
+ awq_enabled: bool = False,
+ flash_attention: bool = False,
+ batch_size: int = 32,
+ notification_email: Optional[str] = None,
+ **kwargs: Any,
+ ) -> None:
+ """
+ Perform bulk translation using the specified model and tokenizer. This method handles the entire translation
+ process including loading the model, processing input data, generating translations, and saving the results.
+
+ Args:
+ model_name (str): Name or path of the translation model.
+ origin (str): Source language ISO code.
+ target (str): Target language ISO code.
+ max_length (int): Maximum length of the tokens (default 512).
+ model_class (str): Class name of the model (default "AutoModelForSeq2SeqLM").
+ tokenizer_class (str): Class name of the tokenizer (default "AutoTokenizer").
+ use_cuda (bool): Whether to use CUDA for model inference (default False).
+ precision (str): Precision for model computation (default "float16").
+ quantization (int): Level of quantization for optimizing model size and speed (default 0).
+ device_map (str | Dict | None): Specific device to use for computation (default "auto").
+ max_memory (Dict): Maximum memory configuration for devices.
+ torchscript (bool, optional): Whether to use a TorchScript-optimized version of the pre-trained language model. Defaults to False.
+ compile (bool, optional): Whether to compile the model before fine-tuning. Defaults to True.
+ awq_enabled (bool): Whether to enable AWQ optimization (default False).
+ flash_attention (bool): Whether to use flash attention optimization (default False).
+ batch_size (int): Number of translations to process simultaneously (default 32).
+ **kwargs: Arbitrary keyword arguments for model and generation configurations.
+ """
+
+ if ":" in model_name:
+ model_revision = model_name.split(":")[1]
+ tokenizer_revision = model_name.split(":")[1]
+ model_name = model_name.split(":")[0]
+ tokenizer_name = model_name
+ else:
+ model_revision = None
+ tokenizer_revision = None
+ tokenizer_name = model_name
+
+ self.model_name = model_name
+ self.tokenizer_name = tokenizer_name
+ self.model_revision = model_revision
+ self.tokenizer_revision = tokenizer_revision
+ self.model_class = model_class
+ self.tokenizer_class = tokenizer_class
+ self.use_cuda = use_cuda
+ self.precision = precision
+ self.quantization = quantization
+ self.device_map = device_map
+ self.max_memory = max_memory
+ self.torchscript = torchscript
+ self.awq_enabled = awq_enabled
+ self.flash_attention = flash_attention
+ self.batch_size = batch_size
+ self.notification_email = notification_email
+ self.compile = compile
+
+ model_args = {k.replace("model_", ""): v for k, v in kwargs.items() if "model_" in k}
+ self.model_args = model_args
+
+ generation_args = {k.replace("generation_", ""): v for k, v in kwargs.items() if "generation_" in k}
+ self.generation_args = generation_args
+
+ self.model, self.tokenizer = self.load_models(
+ model_name=self.model_name,
+ tokenizer_name=self.tokenizer_name,
+ model_revision=self.model_revision,
+ tokenizer_revision=self.tokenizer_revision,
+ model_class=self.model_class,
+ tokenizer_class=self.tokenizer_class,
+ use_cuda=self.use_cuda,
+ precision=self.precision,
+ quantization=self.quantization,
+ device_map=self.device_map,
+ max_memory=self.max_memory,
+ torchscript=self.torchscript,
+ awq_enabled=self.awq_enabled,
+ flash_attention=self.flash_attention,
+ compile=self.compile,
+ **self.model_args,
+ )
+
+ dataset_path = self.input.input_folder
+ output_path = self.output.output_folder
+
+ # Load dataset
+ _dataset = self.load_dataset(dataset_path, origin=origin, target=target, max_length=max_length)
+ if _dataset is None:
+ self.log.error("Failed to load dataset.")
+ return
+ dataset = _dataset[self.origin]
+
+ # Process data in batches
+ for i in range(0, len(dataset), batch_size):
+ batch = dataset[i : i + batch_size]
+ inputs = self.tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
+
+ if next(self.model.parameters()).is_cuda:
+ inputs = {k: v.cuda() for k, v in inputs.items()}
+
+ outputs = self.model.generate(**inputs, **self.generation_args)
+ translations = [self.tokenizer.decode(t, skip_special_tokens=True) for t in outputs]
+
+ self._save_translations(translations, batch, output_path, i)
+ self.done()
+
+ def _save_translations(
+ self, translations: List[str], input_batch: List[str], output_path: str, batch_idx: int
+ ) -> None:
+ r"""
+ Saves the translated texts to a specified output path. This method is called internally by the translate method
+ to persist the translation results.
+
+ Args:
+ translations (List[str]): List of translated texts.
+ input_batch (List[str]): List of original texts that were translated.
+ output_path (str): Path to save the translated texts.
+ batch_idx (int): Index of the current batch (for naming files).
+ """
+
+ data_to_save = [
+ {"input": input_text, "translation": translation}
+ for input_text, translation in zip(input_batch, translations)
+ ]
+ with open(os.path.join(output_path, f"translations-{batch_idx}-{str(uuid.uuid4())}.json"), "w") as f:
+ json.dump(data_to_save, f)
diff --git a/huggingface/translation.py b/geniusrise_text/translation/fine_tune.py
similarity index 72%
rename from huggingface/translation.py
rename to geniusrise_text/translation/fine_tune.py
index 11d5885..477c557 100644
--- a/huggingface/translation.py
+++ b/geniusrise_text/translation/fine_tune.py
@@ -1,18 +1,17 @@
# 🧠 Geniusrise
# Copyright (C) 2023 geniusrise.ai
#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
+# http://www.apache.org/licenses/LICENSE-2.0
#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
- The text content
- The label
-
- ```
-
- ### YAML
- Each document should be a dictionary with 'text' and 'label' keys.
- ```yaml
- - text: "The text content"
- label: "The label"
- ```
-
- ### TSV
- Should contain 'text' and 'label' columns separated by tabs.
-
- ### Excel (.xls, .xlsx)
- Should contain 'text' and 'label' columns.
-
- ### SQLite (.db)
- Should contain a table with 'text' and 'label' columns.
-
- ### Feather
- Should contain 'text' and 'label' columns.
- """
-
- self.data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
- self.max_length = max_length
-
- self.label_to_id = self.model.config.label2id if self.model and self.model.config.label2id else None # type: ignore
-
- def tokenize_function(examples):
- tokenized_data = self.tokenizer(
- examples["text"],
- padding="max_length",
- truncation=True,
- max_length=self.max_length,
- )
- tokenized_data["label"] = [self.label_to_id[label] for label in examples["label"]]
- return tokenized_data
-
- try:
- logging.info(f"Loading dataset from {dataset_path}")
- if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
- # Load dataset saved by Hugging Face datasets library
- return load_from_disk(dataset_path).map(tokenize_function, batched=True)
- else:
- data = []
- for filename in os.listdir(dataset_path):
- filepath = os.path.join(dataset_path, filename)
- if filename.endswith(".jsonl"):
- with open(filepath, "r") as f:
- for line in f:
- example = json.loads(line)
- data.append(example)
-
- elif filename.endswith(".csv"):
- df = pd.read_csv(filepath)
- data.extend(df.to_dict("records"))
-
- elif filename.endswith(".parquet"):
- df = pq.read_table(filepath).to_pandas()
- data.extend(df.to_dict("records"))
-
- elif filename.endswith(".json"):
- with open(filepath, "r") as f:
- json_data = json.load(f)
- data.extend(json_data)
-
- elif filename.endswith(".xml"):
- tree = ET.parse(filepath)
- root = tree.getroot()
- for record in root.findall("record"):
- text = record.find("text").text # type: ignore
- label = record.find("label").text # type: ignore
- data.append({"text": text, "label": label})
-
- elif filename.endswith(".yaml") or filename.endswith(".yml"):
- with open(filepath, "r") as f:
- yaml_data = yaml.safe_load(f)
- data.extend(yaml_data)
-
- elif filename.endswith(".tsv"):
- df = pd.read_csv(filepath, sep="\t")
- data.extend(df.to_dict("records"))
-
- elif filename.endswith((".xls", ".xlsx")):
- df = pd.read_excel(filepath)
- data.extend(df.to_dict("records"))
-
- elif filename.endswith(".db"):
- conn = sqlite3.connect(filepath)
- query = "SELECT text, label FROM dataset_table;"
- df = pd.read_sql_query(query, conn)
- data.extend(df.to_dict("records"))
-
- elif filename.endswith(".feather"):
- df = feather.read_feather(filepath)
- data.extend(df.to_dict("records"))
-
- # Create label_to_id mapping and save it in model config
- unique_labels = (example["label"] for example in data)
- self.label_to_id = {label: i for i, label in enumerate(unique_labels)}
- if self.model:
- if self.model.config.label2id != self.label_to_id:
- self.log.warning("New labels detected, ignore if fine-tuning")
- self.model.config.label2id = self.label_to_id
- self.model.config.id2label = {i: label for label, i in self.label_to_id.items()}
-
- return Dataset.from_pandas(pd.DataFrame(data)).map(tokenize_function, batched=True)
- except Exception as e:
- logging.error(f"Error occurred when loading dataset from {dataset_path}. Error: {e}")
- raise
diff --git a/huggingface/sentiment_analysis.py b/huggingface/sentiment_analysis.py
deleted file mode 100644
index 889cb13..0000000
--- a/huggingface/sentiment_analysis.py
+++ /dev/null
@@ -1,256 +0,0 @@
-# 🧠 Geniusrise
-# Copyright (C) 2023 geniusrise.ai
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see
- The text content
- The label
-
- ```
-
- ### YAML
- Each document should be a dictionary with 'text' and 'label' keys.
- ```yaml
- - text: "The text content"
- label: "The label"
- ```
-
- ### TSV
- Should contain 'text' and 'label' columns separated by tabs.
-
- ### Excel (.xls, .xlsx)
- Should contain 'text' and 'label' columns.
-
- ### SQLite (.db)
- Should contain a table with 'text' and 'label' columns.
-
- ### Feather
- Should contain 'text' and 'label' columns.
- """
- if os.path.isfile(os.path.join(dataset_path, "dataset_info.json")):
- dataset = load_from_disk(dataset_path)
- else:
- data = []
- for filename in os.listdir(dataset_path):
- filepath = os.path.join(dataset_path, filename)
- if filename.endswith(".jsonl"):
- with open(filepath, "r") as f:
- for line in f:
- example = json.loads(line)
- data.append(example)
- elif filename.endswith(".csv"):
- df = pd.read_csv(filepath)
- data.extend(df.to_dict("records"))
- elif filename.endswith(".parquet"):
- df = pq.read_table(filepath).to_pandas()
- data.extend(df.to_dict("records"))
- elif filename.endswith(".json"):
- with open(filepath, "r") as f:
- json_data = json.load(f)
- data.extend(json_data)
- elif filename.endswith(".xml"):
- tree = ET.parse(filepath)
- root = tree.getroot()
- for record in root.findall("record"):
- text = record.find("text").text # type: ignore
- label = record.find("label").text # type: ignore
- data.append({"text": text, "label": label})
- elif filename.endswith(".yaml") or filename.endswith(".yml"):
- with open(filepath, "r") as f:
- yaml_data = yaml.safe_load(f)
- data.extend(yaml_data)
- elif filename.endswith(".tsv"):
- df = pd.read_csv(filepath, sep="\t")
- data.extend(df.to_dict("records"))
- elif filename.endswith((".xls", ".xlsx")):
- df = pd.read_excel(filepath)
- data.extend(df.to_dict("records"))
- elif filename.endswith(".db"):
- conn = sqlite3.connect(filepath)
- query = "SELECT text, label FROM dataset_table;"
- df = pd.read_sql_query(query, conn)
- data.extend(df.to_dict("records"))
- elif filename.endswith(".feather"):
- df = feather.read_feather(filepath)
- data.extend(df.to_dict("records"))
- dataset = Dataset.from_pandas(pd.DataFrame(data))
-
- tokenized_dataset = dataset.map(
- self.prepare_train_features,
- batched=True,
- remove_columns=dataset.column_names,
- )
- return tokenized_dataset
-
- def prepare_train_features(self, examples: Dict[str, Union[str, int]]) -> Dict[str, Union[List[int], int]]:
- """
- Tokenize the examples and prepare the features for training.
-
- Args:
- examples (Dict[str, Union[str, int]]): A dictionary of examples.
-
- Returns:
- Dict[str, Union[List[int], int]]: The processed features.
- """
- if not self.tokenizer:
- raise Exception("No tokenizer found, please call load_models first.")
-
- tokenized_inputs = self.tokenizer(examples["text"], truncation=True, padding=False)
- tokenized_inputs["labels"] = examples["label"]
- return tokenized_inputs
-
- def data_collator(
- self, examples: List[Dict[str, Union[List[int], int]]]
- ) -> Dict[str, Union[torch.Tensor, List[torch.Tensor]]]:
- """
- Customize the data collator.
-
- Args:
- examples (List[Dict[str, Union[List[int], int]]]): The examples to collate.
-
- Returns:
- Dict[str, Union[torch.Tensor, List[torch.Tensor]]]: The collated data.
- """
- return DataCollatorWithPadding(self.tokenizer)(examples)
diff --git a/requirements-dev.txt b/requirements-dev.txt
new file mode 100644
index 0000000..83bf3d1
--- /dev/null
+++ b/requirements-dev.txt
@@ -0,0 +1,334 @@
+absl-py==2.1.0
+accelerate==0.26.1
+aiohttp==3.8.6
+aiosignal==1.3.1
+alembic==1.13.0
+annotated-types==0.6.0
+ansicolors==1.1.8
+anyio==4.1.0
+apache-airflow==2.8.0
+apache-airflow-providers-common-sql==1.9.0
+apache-airflow-providers-docker==3.8.2
+apache-airflow-providers-ftp==3.7.0
+apache-airflow-providers-http==4.8.0
+apache-airflow-providers-imap==3.5.0
+apache-airflow-providers-sqlite==3.6.0
+apache-beam==2.48.0
+apache-flink==1.18.0
+apache-flink-libraries==1.18.0
+apispec==6.3.0
+argcomplete==3.2.1
+argparse-color-formatter==1.2.2.post2
+argparse-manpage==4.4
+asgiref==3.7.2
+async-timeout==4.0.3
+attributedict==0.3.0
+attrs==23.1.0
+auto-gptq==0.6.0
+autoawq==0.1.8
+autocommand==2.2.2
+avro-python3==1.9.2.1
+Babel==2.14.0
+backoff==2.2.1
+bitsandbytes==0.42.0
+black==23.12.1
+bleach==6.0.0
+blessings==1.7
+blinker==1.7.0
+boto3==1.34.23
+botocore==1.34.23
+build==0.10.0
+cachelib==0.9.0
+cachetools==5.3.2
+certifi==2023.11.17
+cffi==1.16.0
+chardet==5.2.0
+charset-normalizer==3.3.2
+cheroot==10.0.0
+CherryPy==18.9.0
+click==8.1.7
+clickclick==20.10.2
+cloudpickle==2.2.1
+codecov==2.1.13
+colorama==0.4.6
+coloredlogs==15.0.1
+colorlog==4.8.0
+colour-runner==0.1.1
+ConfigUpdater==3.2
+connexion==2.14.2
+coverage==7.3.4
+crcmod==1.7
+cron-descriptor==1.4.0
+croniter==2.0.1
+cryptography==41.0.7
+DataProperty==1.0.1
+datasets==2.16.1
+deepdiff==6.7.1
+Deprecated==1.2.14
+dill==0.3.7
+direnv==2020.12.3
+distlib==0.3.8
+dnspython==2.4.2
+docker==7.0.0
+docopt==0.6.2
+docstring-parser==0.15
+docutils==0.20.1
+einops==0.7.0
+email-validator==1.3.1
+emoji==2.7.0
+env-file==2020.12.3
+et-xmlfile==1.1.0
+evaluate==0.4.1
+exceptiongroup==1.2.0
+fastavro==1.4.7
+fasteners==0.19
+fastjsonschema==2.19.1
+filelock==3.13.1
+find-libpython==0.3.0
+flake8==6.1.0
+flash-attn==2.4.2
+Flask==2.2.5
+Flask-AppBuilder==4.3.10
+Flask-Babel==2.0.0
+Flask-Caching==2.1.0
+Flask-JWT-Extended==4.6.0
+Flask-Limiter==3.5.0
+Flask-Login==0.6.3
+Flask-Session==0.5.0
+Flask-SQLAlchemy==2.5.1
+Flask-WTF==1.2.1
+frozenlist==1.4.1
+fsspec==2023.10.0
+gekko==1.0.6
+geniusrise==0.0.33
+google-auth==2.17.3
+google-re2==1.1
+googleapis-common-protos==1.62.0
+GPUtil==1.4.0
+graphviz==0.20.1
+greenlet==3.0.2
+grpcio==1.60.0
+gunicorn==21.2.0
+h11==0.14.0
+hdfs==2.7.2
+httpcore==0.16.3
+httplib2==0.20.4
+httpx==0.23.3
+huggingface-hub==0.20.2
+humanfriendly==10.0
+idna==3.6
+importlib-metadata==6.11.0
+importlib-resources==6.1.1
+inflect==7.0.0
+inflection==0.5.1
+iniconfig==2.0.0
+inspecta==0.1.3
+itsdangerous==2.1.2
+jaraco.classes==3.3.0
+jaraco.collections==5.0.0
+jaraco.context==4.3.0
+jaraco.functools==4.0.0
+jaraco.text==3.12.0
+jeepney==0.8.0
+Jinja2==3.1.2
+jmespath==0.10.0
+joblib==1.3.2
+jsonlines==4.0.0
+jsonpickle==3.0.1
+jsonschema==4.20.0
+jsonschema-specifications==2023.11.2
+jupyter_core==5.7.1
+kafka-python==2.0.2
+keyring==24.2.0
+kubernetes==28.1.0
+lazy-object-proxy==1.10.0
+limits==3.7.0
+linkify-it-py==2.0.2
+lm_eval==0.4.0
+lockfile==0.12.2
+lxml==5.1.0
+Mako==1.3.0
+Markdown==3.5.1
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+marshmallow==3.20.1
+marshmallow-oneofschema==3.0.1
+marshmallow-sqlalchemy==0.26.1
+mbstrdecoder==1.1.3
+mccabe==0.7.0
+mdit-py-plugins==0.4.0
+mdurl==0.1.2
+more-itertools==10.1.0
+mpmath==1.3.0
+multidict==6.0.4
+multiprocess==0.70.15
+mypy==1.8.0
+mypy-extensions==1.0.0
+nbformat==5.9.2
+networkx==3.2.1
+ninja==1.11.1.1
+nltk==3.8.1
+numexpr==2.8.8
+numpy==1.26.3
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.18.1
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu12==12.1.105
+oauthlib==3.2.2
+objsize==0.6.1
+openpyxl==3.1.2
+opentelemetry-api==1.22.0
+opentelemetry-exporter-otlp==1.22.0
+opentelemetry-exporter-otlp-proto-common==1.22.0
+opentelemetry-exporter-otlp-proto-grpc==1.22.0
+opentelemetry-exporter-otlp-proto-http==1.22.0
+opentelemetry-proto==1.22.0
+opentelemetry-sdk==1.22.0
+opentelemetry-semantic-conventions==0.43b0
+ordered-set==4.1.0
+orjson==3.9.7
+packaging==23.2
+pandas==1.3.5
+pathspec==0.12.1
+pathvalidate==3.2.0
+peft==0.7.1
+pemja==0.3.0
+pendulum==2.1.2
+pillow==10.2.0
+pip-autoremove==0.10.0
+pipdeptree==2.13.1
+pkginfo==1.9.6
+platformdirs==4.1.0
+pluggy==1.3.0
+portalocker==2.8.2
+portend==3.2.0
+prettytable==3.9.0
+prison==0.2.1
+prometheus-client==0.19.0
+proto-plus==1.22.3
+protobuf==4.23.4
+psutil==5.9.6
+psycopg2==2.9.9
+py4j==0.10.9.7
+pyarrow==8.0.0
+pyarrow-hotfix==0.6
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pybind11==2.11.1
+pycodestyle==2.11.0
+pycparser==2.21
+pydantic==2.5.2
+pydantic_core==2.14.5
+pydot==1.4.2
+pyflakes==3.1.0
+Pygments==2.17.2
+PyJWT==2.8.0
+pymongo==3.13.0
+pyparsing==3.1.1
+pyproject-api==1.6.1
+pyproject_hooks==1.0.0
+pyspark==3.5.0
+pytablewriter==1.2.0
+pytest==7.4.3
+pytest-asyncio==0.21.1
+python-daemon==3.0.1
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+python-nvd3==0.15.0
+python-slugify==8.0.1
+pytz==2023.3.post1
+pytzdata==2020.1
+PyYAML==6.0.1
+readme-renderer==40.0
+redis==5.0.1
+referencing==0.32.0
+regex==2023.8.8
+requests==2.31.0
+requests-oauthlib==1.3.1
+requests-toolbelt==1.0.0
+responses==0.18.0
+retrying==1.3.4
+rfc3339-validator==0.1.4
+rfc3986==1.5.0
+rich==13.7.0
+rich-argparse==1.4.0
+rootpath==0.1.1
+rouge==1.0.1
+rouge-score==0.1.2
+rpds-py==0.13.2
+rsa==4.9
+s3transfer==0.10.0
+sacrebleu==2.4.0
+sacremoses==0.1.1
+safetensors==0.4.1
+scikit-learn==1.4.0
+scipy==1.12.0
+SecretStorage==3.3.3
+sentence-transformers==2.2.2
+sentencepiece==0.1.99
+setproctitle==1.3.3
+shortuuid==1.0.11
+shtab==1.6.5
+six==1.16.0
+sniffio==1.3.0
+SQLAlchemy==1.4.50
+SQLAlchemy-JSONField==1.0.2
+SQLAlchemy-Utils==0.41.1
+sqlitedict==2.1.0
+sqlparse==0.4.4
+streamz==0.6.4
+sympy==1.12
+tabledata==1.3.3
+tabulate==0.9.0
+tcolorpy==0.1.4
+tempora==5.5.0
+tenacity==8.2.3
+termcolor==2.4.0
+text-unidecode==1.3
+texttable==1.7.0
+threadpoolctl==3.2.0
+tokenizers==0.15.0
+toml==0.10.2
+tomli==2.0.1
+toolz==0.12.0
+torch==2.1.2
+torchvision==0.16.2
+tornado==6.3.3
+tox==4.12.1
+tqdm==4.66.1
+tqdm-multiprocess==0.0.11
+traitlets==5.14.1
+transformers==4.36.2
+triton==2.1.0
+trl==0.7.10
+twine==4.0.2
+typepy==1.3.2
+typing_extensions==4.9.0
+tyro==0.6.6
+tzdata==2023.3
+uc-micro-py==1.0.2
+unicodecsv==0.14.1
+universal-pathlib==0.1.4
+urllib3==1.26.18
+values==2020.12.3
+virtualenv==20.25.0
+wcwidth==0.2.6
+webencodings==0.5.1
+websocket-client==1.6.1
+Werkzeug==2.2.3
+wrapt==1.16.0
+WTForms==3.1.1
+xxhash==3.4.1
+yarl==1.9.4
+zc.lockfile==3.0.post1
+zict==3.0.0
+zipp==3.17.0
+zstandard==0.21.0
diff --git a/requirements.txt b/requirements.txt
index ef38b0f..69e3aa2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,147 +1,35 @@
-absl-py==1.4.0
-accelerate==0.22.0
-aiohttp==3.8.5
-aiosignal==1.3.1
-annotated-types==0.5.0
-ansicolors==1.1.8
-argparse-color-formatter==1.2.2.post2
-async-timeout==4.0.3
-attrs==23.1.0
-bleach==6.0.0
-boto3==1.28.25
-botocore==1.31.25
-build==0.10.0
-cachetools==5.3.1
-certifi==2023.7.22
-cffi==1.15.1
-charset-normalizer==3.2.0
-click==8.1.7
-cmake==3.27.2
-colorama==0.4.6
-colorlog==6.7.0
-coverage==7.3.0
-cryptography==41.0.3
-datasets==2.14.4
-dill==0.3.7
-direnv==2020.12.3
-docutils==0.20.1
-emoji==2.7.0
-env-file==2020.12.3
-et-xmlfile==1.1.0
-evaluate==0.4.0
-exceptiongroup==1.1.2
-filelock==3.12.2
-flake8==6.1.0
-frozenlist==1.4.0
-fsspec==2023.6.0
-geniusrise==0.0.1
-google-auth==2.17.3
-huggingface-hub==0.16.4
-idna==3.4
-importlib-metadata==6.8.0
-iniconfig==2.0.0
-jaraco.classes==3.3.0
-jeepney==0.8.0
-Jinja2==3.1.2
-jmespath==0.10.0
-joblib==1.3.2
-jsonpickle==3.0.1
-kafka-python==2.0.2
-keyring==24.2.0
-kubernetes==27.2.0
-lit==16.0.6
-lxml==4.9.3
-markdown-it-py==3.0.0
-MarkupSafe==2.1.3
-mccabe==0.7.0
-mdurl==0.1.2
-more-itertools==10.1.0
-mpmath==1.3.0
-multidict==6.0.4
-multiprocess==0.70.15
-mypy==1.5.0
-mypy-extensions==1.0.0
-networkx==3.1
-nltk==3.8.1
-numpy==1.25.2
-nvidia-cublas-cu11==11.10.3.66
-nvidia-cuda-cupti-cu11==11.7.101
-nvidia-cuda-nvrtc-cu11==11.7.99
-nvidia-cuda-runtime-cu11==11.7.99
-nvidia-cudnn-cu11==8.5.0.96
-nvidia-cufft-cu11==10.9.0.58
-nvidia-curand-cu11==10.2.10.91
-nvidia-cusolver-cu11==11.4.0.1
-nvidia-cusparse-cu11==11.7.4.91
-nvidia-nccl-cu11==2.14.3
-nvidia-nvtx-cu11==11.7.91
-oauthlib==3.2.2
-openpyxl==3.1.2
-packaging==23.1
-pandas==2.0.3
-pkginfo==1.9.6
-pluggy==1.2.0
-portalocker==2.7.0
-prettytable==3.8.0
-psutil==5.9.5
-psycopg2==2.9.7
-pyarrow==13.0.0
-pyasn1==0.5.0
-pyasn1-modules==0.3.0
-pycodestyle==2.11.0
-pycparser==2.21
-pydantic==2.1.1
-pydantic_core==2.4.0
-pyflakes==3.1.0
-Pygments==2.16.1
-pyproject_hooks==1.0.0
-pytest==7.4.0
-pytest-cov==4.1.0
-python-dateutil==2.8.2
-pytz==2023.3
-PyYAML==6.0.1
-readme-renderer==40.0
-redis==4.6.0
-regex==2023.8.8
-requests==2.31.0
-requests-oauthlib==1.3.1
-requests-toolbelt==1.0.0
-responses==0.18.0
-retrying==1.3.4
-rfc3986==2.0.0
-rich==13.5.2
-rich-argparse==1.3.0
+geniusrise==0.1.7
+torch==2.1.2
+accelerate==0.27.2
+transformers==4.38.1
+evaluate==0.4.1
+datasets==2.16.1
+autoawq==0.2.2
+auto-gptq==0.7.0
+bitsandbytes==0.42.0
+CherryPy==18.9.0
+flash-attn==2.4.2
+peft==0.7.1
+gunicorn==21.2.0
+trl==0.7.10
+sentence_transformers==2.2.2
+nbformat==5.9.2
+sacremoses==0.1.1
+sacrebleu==2.4.0
rouge-score==0.1.2
-rsa==4.9
-s3transfer==0.6.1
-sacrebleu==2.3.1
-sacremoses==0.0.53
-safetensors==0.3.3
-scikit-learn==1.3.0
-scipy==1.11.2
-SecretStorage==3.3.3
+rouge==1.0.1
+openpyxl==3.1.2
sentencepiece==0.1.99
-shortuuid==1.0.11
-six==1.16.0
-sympy==1.12
-tabulate==0.9.0
-termcolor==2.3.0
-threadpoolctl==3.2.0
-tokenizers==0.13.3
-tomli==2.0.1
-torch==2.0.1
-tqdm==4.66.1
-transformers==4.32.0
-triton==2.0.0
-twine==4.0.2
-types-PyYAML==6.0.12.11
-typing_extensions==4.7.1
-tzdata==2023.3
-urllib3==1.26.16
-values==2020.12.3
-wcwidth==0.2.6
-webencodings==0.5.1
-websocket-client==1.6.1
-xxhash==3.3.0
-yarl==1.9.2
-zipp==3.16.2
+safetensors==0.4.1
+einops==0.7.0
+packaging==23.2
+optimum==1.16.2
+textwrap3==0.9.2
+nltk==3.8.1
+pyarrow>=9.0.0
+absl-py==2.1.0
+pandas==1.3.5
+multiprocess==0.70.15
+optimum==1.16.2
+vllm==0.3.2
+llama_cpp_python==0.2.50
diff --git a/setup.cfg b/setup.cfg
index c092ec3..54e5d5b 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
[flake8]
exclude = */__init__.py,migrations/*
-ignore = E111, E114, E121, E131, W503, F405, F403, E126, E501, F841, E124, E251, E203, C419, B008, W505, N805, A003, B006, A002, B009, N817, A001
+ignore = E111, E114, E121, E131, W503, F405, F403, E126, E501, F841, E124, E251, E203, C419, B008, W505, N805, A003, B006, A002, B009, N817, A001, N802, N806, E741, C416, C400
max-line-length = 120
max-doc-length = 120
show-source = true
@@ -21,3 +21,7 @@ disallow_untyped_calls = False
[tool:pytest]
addopts = -p no:warnings
ignore = tests
+
+[tool.ruff]
+indent-width = 4
+line-length = 120
diff --git a/setup.py b/setup.py
index 3509e5a..3b8e77d 100644
--- a/setup.py
+++ b/setup.py
@@ -3,21 +3,42 @@
with open("README.md", "r", encoding="utf-8") as fh:
long_description = fh.read()
+with open("requirements.txt", "r", encoding="utf-8") as f:
+ requirements = f.read().splitlines()
+
setup(
- name="geniusrise-huggingface",
- version="0.1.0",
+ name="geniusrise-text",
+ version="0.1.12",
packages=find_packages(exclude=["tests", "tests.*"]),
- install_requires=[],
+ install_requires=requirements,
python_requires=">=3.10",
author="ixaxaar",
author_email="ixaxaar@geniusrise.ai",
description="Huggingface bolts for geniusrise",
long_description=long_description,
long_description_content_type="text/markdown",
- url="https://github.com/ixaxaar/huggingface-bolts",
+ url="https://github.com/geniusrise/geniusrise-text",
classifiers=[
- "Programming Language :: Python :: 3",
- "License :: OSI Approved :: MIT License",
+ "Development Status :: 3 - Alpha",
+ "Intended Audience :: Developers",
+ "Topic :: Software Development :: Build Tools",
+ "License :: OSI Approved :: Apache Software License",
+ "Programming Language :: Python :: 3.10",
+ "Programming Language :: Python :: 3.11",
+ "Programming Language :: Python :: 3.12",
"Operating System :: OS Independent",
],
+ keywords="mlops, llm, geniusrise, machine learning, data processing",
+ project_urls={
+ "Bug Reports": "https://github.com/geniusrise/geniusrise-text/issues",
+ "Source": "https://github.com/geniusrise/geniusrise-text",
+ "Documentation": "https://docs.geniusrise.ai/",
+ },
+ package_data={
+ "geniusrise": [],
+ },
+ extras_require={
+ "dev": ["check-manifest"],
+ "test": ["coverage"],
+ },
)
diff --git a/tests/__init__.py b/tests/__init__.py
deleted file mode 100644
index 7e7658f..0000000
--- a/tests/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# 🧠 Geniusrise
-# Copyright (C) 2023 geniusrise.ai
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see