@@ -573,10 +589,10 @@ def main():
labeller = st.radio(
"šØāš¼ Labeller",
[
+ "Golden (only for benchmarking)",
"Open-source / Custom LLM",
"API LLM",
"Human",
- "Golden (only for benchmarking)",
],
help="Select the type of labeller to use for data annotation",
)
@@ -592,12 +608,29 @@ def main():
help="Cost paid to human annotators per example",
)
elif labeller == "custom_llm":
- model_checkpoint = st.text_input(
+ labeller_checkpoint = st.text_input(
"š¤ Model checkpoint from HuggingFace",
- value="Qwen/QwQ-32B",
+ value="Qwen/Qwen3-32B",
help="HuggingFace model ID for the custom LLM",
)
elif labeller == "api_llm":
+ # Add data privacy disclaimer
+ st.markdown(
+ """
+
+
+ ā ļø Data Privacy Notice
+
+
+ Important: When using API-based labellers (OpenAI, Anthropic, etc.), your dataset will be sent to external services for processing.
+ Please ensure you have the necessary permissions and that your data complies with the respective service providers' terms of use and privacy policies.
+ Consider using local/custom models if your data contains sensitive or proprietary information.
+
+
+ """,
+ unsafe_allow_html=True,
+ )
+
col1, col2 = st.columns(2)
with col1:
provider = st.radio(
@@ -664,7 +697,7 @@ def main():
# Dataset input
dataset = st.text_input(
"š Dataset or path to data",
- value="SpeedOfMagic/gigaword_tiny",
+ value="Yale-LILY/aeslc",
help="HuggingFace dataset ID or local path to dataset",
)
@@ -788,13 +821,13 @@ def main():
with col1:
input_field = st.text_input(
"š„ Input field name",
- value="document",
+ value="email_body",
help="Name of the field containing input text in the dataset",
)
with col2:
reference_field = st.text_input(
"š¤ Reference field name",
- value="summary",
+ value="subject_line",
help="Name of the field containing reference output in the dataset",
)
@@ -822,7 +855,7 @@ def main():
with col1:
model_checkpoint = st.text_input(
"š¤ Model checkpoint",
- value="Qwen/Qwen2.5-1.5B-Instruct",
+ value="Qwen/Qwen3-1.7B",
help="HuggingFace model ID for generation",
)
@@ -1142,9 +1175,11 @@ def main():
# Check again if an experiment is already running
if check_experiment_status():
- st.error("Another experiment is already running. Please wait for it to complete.")
+ st.error(
+ "Another experiment is already running. Please wait for it to complete."
+ )
st.stop()
-
+
# Create a progress container outside of any spinner
progress_container = st.container()
with progress_container:
@@ -1157,7 +1192,7 @@ def main():
""",
unsafe_allow_html=True,
)
-
+
# Add progress tracking elements
progress_bar = st.progress(0)
progress_status = st.empty()
@@ -1218,14 +1253,16 @@ def main():
# Process uploaded datasets if available
if train_dataset_path is not None:
config = process_uploaded_datasets(
- train_dataset_path=train_dataset_path,
- test_dataset_path=test_dataset_path,
- output_dir=output_dir,
- config=config,
- status=status
+ train_dataset_path=train_dataset_path,
+ test_dataset_path=test_dataset_path,
+ output_dir=output_dir,
+ config=config,
+ status=status,
)
else:
- st.error("No training dataset uploaded. Please upload a training dataset.")
+ st.error(
+ "No training dataset uploaded. Please upload a training dataset."
+ )
st.stop()
else:
# Using standard dataset (HuggingFace or local path)
@@ -1266,7 +1303,9 @@ def main():
] = price_input_per_example
elif labeller == "custom_llm":
if "model_checkpoint" in locals():
- config["labeller"]["model"]["checkpoint"] = model_checkpoint
+ config["labeller"]["model"][
+ "checkpoint"
+ ] = labeller_checkpoint
elif labeller == "api_llm":
config["labeller"]["api_key"] = api_key
config["labeller"]["provider"] = provider
@@ -1370,8 +1409,7 @@ def main():
)
# Create progress tracking callback
progress_update_callback = create_progress_tracker(
- progress_container,
- num_iterations
+ progress_container, num_iterations
)
# Clear the status before running the experiment
@@ -1380,13 +1418,12 @@ def main():
# Run the active learning experiment with progress tracking
with st.spinner("Running active learning experiment..."):
run_active_learning_with_progress(
- config,
- progress_callback=progress_update_callback
+ config, progress_callback=progress_update_callback
)
# Update status to completed
update_experiment_status(STATUS_COMPLETED)
-
+
# Show success indicators after completion
st.balloons()
@@ -1401,24 +1438,24 @@ def main():
nav_col1, nav_col2, nav_col3 = st.columns(3)
with nav_col1:
if st.button("š View Metrics", use_container_width=True):
- st.switch_page("1_Metrics.py")
+ st.switch_page("1_Metrics")
with nav_col2:
if st.button(
"š·ļø View Labeled Examples", use_container_width=True
):
- st.switch_page("2_Labeled_examples.py")
+ st.switch_page("2_Labeled_examples")
with nav_col3:
if st.button(
"š©āšØ Annotate Examples", use_container_width=True
):
- st.switch_page("3_Annotation.py")
+ st.switch_page("3_Annotation")
except Exception as e:
# Update status to failed
update_experiment_status(STATUS_FAILED)
- st.error(f"An error occurred: {str(e)}")
- import traceback
-
- st.code(traceback.format_exc(), language="python")
+ # st.error(f"An error occurred: {str(e)}")
+ # import sys, pdb
+ # exc_type, exc_value, exc_traceback = sys.exc_info()
+ # pdb.post_mortem(exc_traceback)
elif is_valid_required_performance:
st.error("ā You didn't fill one of the required arguments.")
diff --git a/pyproject.toml b/pyproject.toml
index 125ac23..bcfe10d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "atgen"
-version = "0.0.0"
+version = "0.1.0"
authors = [
{ name = "List of contributors: https://github.com/Aktsvigun/atgen/graphs/contributors", email = "artemshelmanov@gmail.com" },
]
@@ -12,7 +12,7 @@ description = "Toolkit for Active Learning in Generative Tasks"
readme = "README.md"
keywords = ["NLP", "deep learning", "transformer", "pytorch", "peft", "inference", "active learning"]
-license = {file = "LICENSE.md"}
+license-files = ["LICENSE.md"]
requires-python = ">=3.10"
dynamic = ["dependencies"]
diff --git a/requirements.txt b/requirements.txt
index 3742987..216563c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+alignscore-SpeedOfMagic
accelerate==1.5.2
anthropic==0.49.0
benepar==0.2.0
@@ -27,7 +28,7 @@ spacy==3.7.5
streamlit==1.37.0
streamlit-authenticator==0.4.2
tabulate==0.9.0
-transformers==4.49.0
+transformers==4.52.4
trl==0.15.2
torchmetrics==1.4.1
unsloth==2025.3.17
diff --git a/src/atgen/labellers/api_labellers/openai_labeller.py b/src/atgen/labellers/api_labellers/openai_labeller.py
index 529cfbc..0c0530e 100644
--- a/src/atgen/labellers/api_labellers/openai_labeller.py
+++ b/src/atgen/labellers/api_labellers/openai_labeller.py
@@ -10,7 +10,7 @@
from shutil import rmtree
from ..base_labeller import BaseLabeler
-from ...utils.constants import DEEPSEEK_R1_END_REASONING_TOKEN, MESSAGES_COLUMN_NAME
+from ...utils.constants import REASONING_END_TOKEN, MESSAGES_COLUMN_NAME
log = logging.getLogger()
@@ -126,8 +126,8 @@ def _sync_call(self, dataset: Dataset) -> Dataset:
# Remove thinking tokens from DeepSeek-R1
if "deepseek-r1" in self.config.parameters.model:
for i, annotation in enumerate(annotations):
- annotations[i] = DEEPSEEK_R1_END_REASONING_TOKEN.join(
- annotation.split(DEEPSEEK_R1_END_REASONING_TOKEN)[1:]
+ annotations[i] = REASONING_END_TOKEN.join(
+ annotation.split(REASONING_END_TOKEN)[1:]
).strip()
# Add the annotations as a new column in the dataset
diff --git a/src/atgen/labellers/golden_labeller.py b/src/atgen/labellers/golden_labeller.py
index 6d2fd45..cf02190 100644
--- a/src/atgen/labellers/golden_labeller.py
+++ b/src/atgen/labellers/golden_labeller.py
@@ -1,3 +1,5 @@
+from typing import Union
+from omegaconf import DictConfig, ListConfig
from datasets import Dataset
from .base_labeller import BaseLabeler
@@ -6,7 +8,36 @@
# Labels rows by using their label (should be used when evaluating strategies)
class GoldenLabeler(BaseLabeler):
def __call__(self, dataset: Dataset) -> Dataset:
- assert (
- self.output_column_name in dataset.column_names
- ), "Column with labels was not found in dataset"
+ _check_output_column_in_dataset(
+ column_names=self.output_column_name, dataset=dataset
+ )
return dataset
+
+
+def _check_output_column_in_dataset(
+ column_names: Union[
+ DictConfig, ListConfig, dict[str, Union[str, list[str]]], list[str], str
+ ],
+ dataset: Union[Dataset, dict[str, Union[str, dict[str, str]]]],
+) -> None:
+ if isinstance(column_names, str):
+ assert (
+ column_names in dataset.column_names
+ ), f"Column {column_names} with labels was not found in dataset"
+ elif isinstance(column_names, (ListConfig, list)):
+ if isinstance(dataset, Dataset):
+ assert (
+ column_names[0] in dataset.column_names
+ ), f"Column {column_names[0]} with labels was not found in dataset"
+ dataset = dataset[0]
+ else:
+ assert (
+ column_names[0] in dataset.keys()
+ ), f"Column {column_names[0]} with labels was not found in dataset"
+ if len(column_names) > 1:
+ _check_output_column_in_dataset(
+ column_names=column_names[1:], dataset=dataset[column_names[0]]
+ )
+ elif isinstance(column_names, (DictConfig, dict)):
+ for purpose, column_name in column_names.items():
+ _check_output_column_in_dataset(column_names=column_name, dataset=dataset)
diff --git a/src/atgen/metrics/__init__.py b/src/atgen/metrics/__init__.py
index 302149a..8f9078e 100644
--- a/src/atgen/metrics/__init__.py
+++ b/src/atgen/metrics/__init__.py
@@ -1,12 +1,122 @@
-AVAILABLE_METRICS = [
- "rouge1",
- "rouge2",
- "rougeL",
- "rougeLsum",
- "bleu",
- "sacrebleu",
- "BARTScore-hr",
- "BARTScore-sh",
- "alignscore",
- "BigBenchHardScore",
+# Base classes
+from .base import BaseMetric, MetricConfig
+
+# Lexical metrics
+from .lexical import BleuMetric, RougeMetric
+
+# Semantic metrics
+from .semantic import BartScoreMetric, AlignScoreMetric, SentBertMetric
+
+# Linguistic metrics
+from .linguistic import ColaMetric
+
+# LLM-based metrics
+from .llm_based import (
+ EvaluationLLM,
+ BaseDeepEvalMetric,
+ DeepEvalAnswerRelevancyMetric,
+ DeepEvalFaithfulnessMetric,
+ DeepEvalSummarizationMetric,
+ DeepEvalPromptAlignmentMetric,
+ BigBenchHardMetric,
+)
+
+# Factory system
+from .factory import (
+ MetricsFactory,
+ MetricsConfig,
+ get_metric_categories,
+ get_metric_requirements,
+)
+
+# Compute system
+from .compute_metrics import (
+ compute_metrics,
+ compute_metrics_from_config,
+ get_default_config,
+ get_comprehensive_config,
+ get_deepeval_config,
+)
+
+
+
+
+
+# Version
+__version__ = "2.0.0"
+
+# Main exports
+__all__ = [
+ # Base classes
+ "BaseMetric",
+ "MetricConfig",
+
+ # Lexical metrics
+ "BleuMetric",
+ "RougeMetric",
+
+ # Semantic metrics
+ "BartScoreMetric",
+ "AlignScoreMetric",
+ "SentBertMetric",
+
+ # Linguistic metrics
+ "ColaMetric",
+
+ # LLM-based metrics
+ "EvaluationLLM",
+ "BaseDeepEvalMetric",
+ "DeepEvalAnswerRelevancyMetric",
+ "DeepEvalFaithfulnessMetric",
+ "DeepEvalSummarizationMetric",
+ "DeepEvalPromptAlignmentMetric",
+ "BigBenchHardMetric",
+
+ # Factory system
+ "MetricsFactory",
+ "MetricsConfig",
+ "get_metric_categories",
+ "get_metric_requirements",
+
+ # New compute system
+ "compute_metrics",
+ "compute_metrics_from_config",
+ "get_default_config",
+ "get_comprehensive_config",
+ "get_deepeval_config",
+
+ # Helper functions
+ "get_available_metrics",
+ "get_all_possible_metric_keys",
+ "create_metric",
+ "create_metrics_from_config",
+
+ "AVAILABLE_METRICS",
+ "METRICS",
]
+
+
+def get_available_metrics():
+ """Get all available metric names."""
+ return MetricsFactory.get_available_metrics()
+
+
+def get_all_possible_metric_keys():
+ """Get all possible metric keys that can be returned by compute_metrics."""
+ return MetricsFactory.get_all_possible_metric_keys()
+
+
+def create_metric(name: str, config: MetricConfig = None):
+ """Create a metric by name."""
+ return MetricsFactory.create_metric(name, config)
+
+
+def create_metrics_from_config(config):
+ """Create metrics from configuration."""
+ return MetricsFactory.create_from_config(config)
+
+# For backward compatibility, still provide the basic metric names
+METRICS = MetricsFactory.get_available_metrics()
+
+# But AVAILABLE_METRICS should include all possible keys for performance checking
+AVAILABLE_METRICS = MetricsFactory.get_all_possible_metric_keys()
diff --git a/src/atgen/metrics/bart_score.py b/src/atgen/metrics/bart_score.py
index 1e8f808..6133693 100644
--- a/src/atgen/metrics/bart_score.py
+++ b/src/atgen/metrics/bart_score.py
@@ -42,7 +42,7 @@ def load(self, path=None):
def score(self, srcs, tgts, batch_size=4):
"""Score a batch of examples"""
score_list = []
- for i in tqdm(range(0, len(srcs), batch_size), desc="Calculating BARTScore..."):
+ for i in range(0, len(srcs), batch_size):
src_list = srcs[i : i + batch_size]
tgt_list = tgts[i : i + batch_size]
try:
diff --git a/src/atgen/metrics/base/__init__.py b/src/atgen/metrics/base/__init__.py
new file mode 100644
index 0000000..a0f1510
--- /dev/null
+++ b/src/atgen/metrics/base/__init__.py
@@ -0,0 +1,8 @@
+"""Base classes and configuration for metrics."""
+
+from .base_metric import BaseMetric, MetricConfig
+
+__all__ = [
+ "BaseMetric",
+ "MetricConfig",
+]
\ No newline at end of file
diff --git a/src/atgen/metrics/base/base_metric.py b/src/atgen/metrics/base/base_metric.py
new file mode 100644
index 0000000..8c1fb84
--- /dev/null
+++ b/src/atgen/metrics/base/base_metric.py
@@ -0,0 +1,99 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Dict, Any
+import logging
+from dataclasses import dataclass, field
+
+
+@dataclass
+class MetricConfig:
+<<<<<<< HEAD
+=======
+ """Base configuration for metrics."""
+ # General parameters
+>>>>>>> a24e0f2 (removed dependencies chek)
+ batch_size: int = 32
+ device: str = "cuda"
+ cache_dir: str = "cache"
+ aggregate: bool = True
+
+<<<<<<< HEAD
+ checkpoint: Optional[str] = None
+ model_name: Optional[str] = None
+
+ # API configuration parameters
+ provider: Optional[str] = None
+=======
+ # Model-specific parameters (for local models)
+ checkpoint: Optional[str] = None
+ model_name: Optional[str] = None
+
+ # API-based parameters (for LLM-based metrics)
+>>>>>>> a24e0f2 (removed dependencies chek)
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ model: Optional[str] = None
+
+<<<<<<< HEAD
+=======
+ # DeepEval specific parameters
+>>>>>>> a24e0f2 (removed dependencies chek)
+ threshold: float = 0.5
+ include_reason: bool = False
+ strict_mode: bool = False
+ async_mode: bool = True
+ verbose_mode: bool = False
+ truths_extraction_limit: Optional[int] = None
+<<<<<<< HEAD
+=======
+
+ # BigBenchHard specific parameters
+>>>>>>> a24e0f2 (removed dependencies chek)
+ benchmark_params: Dict[str, Any] = field(default_factory=dict)
+ generation_params: Dict[str, Any] = field(default_factory=dict)
+
+
+class BaseMetric(ABC):
+ def __init__(self, config: Optional[MetricConfig] = None):
+ self.config = config or MetricConfig()
+ self.logger = logging.getLogger(self.__class__.__name__)
+<<<<<<< HEAD
+
+ self._validate_config()
+
+ @property
+ def name(self) -> str:
+ return self.__class__.__name__.replace("Metric", "").lower()
+
+ def _validate_config(self):
+=======
+ self._is_available = None
+ self._model = None
+
+ # Validate configuration
+ self._validate_config()
+
+ # Check dependencies
+ if not self.is_available():
+ self.logger.warning(f"{self.__class__.__name__} dependencies not available")
+
+ @property
+ def name(self) -> str:
+ """Return the metric name (defaults to class name)."""
+ return self.__class__.__name__.replace("Metric", "").lower()
+
+ def _validate_config(self):
+ """Validate configuration - can be overridden by subclasses."""
+>>>>>>> a24e0f2 (removed dependencies chek)
+ pass
+
+ @abstractmethod
+ def calculate(self, predictions, references, original_texts):
+ pass
+<<<<<<< HEAD
+
+ def is_available(self) -> bool:
+ """Check if the metric is available for use."""
+ return True
+=======
+>>>>>>> a24e0f2 (removed dependencies chek)
+
\ No newline at end of file
diff --git a/src/atgen/metrics/base/config.py b/src/atgen/metrics/base/config.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/atgen/metrics/big_bench_hard_metric.py b/src/atgen/metrics/big_bench_hard_metric.py
deleted file mode 100644
index 1ec5227..0000000
--- a/src/atgen/metrics/big_bench_hard_metric.py
+++ /dev/null
@@ -1,120 +0,0 @@
-from deepeval.benchmarks import BigBenchHard
-from deepeval.benchmarks.big_bench_hard.template import BigBenchHardTemplate
-from deepeval.models.base_model import DeepEvalBaseLLM
-from transformers import (
- GenerationMixin,
- PreTrainedTokenizerBase,
-)
-
-"""
-https://arxiv.org/abs/2210.09261v1
-
-For benchmark_params, refer to https://docs.confident-ai.com/docs/benchmarks-big-bench-hard
-"""
-
-
-class BigBenchHardModel(DeepEvalBaseLLM):
- def __init__(
- self,
- model: GenerationMixin,
- tokenizer: PreTrainedTokenizerBase,
- device: str,
- model_name: str = "Model",
- **generation_kwargs,
- ) -> None:
- self.model = model
- self.tokenizer = tokenizer
- self.device = device
- self.model_name = model_name
- self.generation_kwargs = generation_kwargs
-
- def load_model(self) -> GenerationMixin:
- return self.model
-
- def generate(self, prompt: str, **kwargs) -> str:
- model = self.load_model()
- model.to(self.device)
- model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device)
- generated_ids = model.generate(**model_inputs, **self.generation_kwargs)
- return self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
- async def a_generate(self, prompt: str, **kwargs) -> str:
- return self.generate(prompt)
-
- def batch_generate(self, prompts: list[str], **kwargs) -> list[str]:
- model = self.load_model()
- model.to(self.device)
- model_inputs = self.tokenizer(prompts, return_tensors="pt").to(self.device)
- generated_ids = model.generate(**model_inputs, **self.generation_kwargs)
- return self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-
- def get_model_name(self):
- return self.model_name
-
-
-"""
-The main branch for deepeval seems to contain broken code, as "NumberModel" is not defined:
-https://github.com/confident-ai/deepeval/blob/main/deepeval/benchmarks/big_bench_hard/big_bench_hard.py#L153
-
-This child class fixes batch prediction.
-"""
-
-
-class BigBenchHardFixed(BigBenchHard):
- def batch_predict(self, model, task, goldens):
- prompts = []
- for golden in goldens:
- prompt: dict = BigBenchHardTemplate.generate_output(
- input=golden.input,
- task=task,
- n_shots=self.n_shots,
- enable_cot=self.enable_cot,
- )
- prompts.append(prompt)
-
- # Enforced model generation
- prompts = [
- prompt + "Make sure to output only the numerical answer."
- for prompt in prompts
- ]
- predictions = model.batch_generate(prompts)
- predictions = [str(pred) for pred in predictions]
-
- if len(predictions) is not len(goldens):
- raise ValueError(
- "Custom `batch_generate` method did not return the same "
- "number of generations as the number of prompts."
- )
-
- res = []
- for i in range(len(predictions)):
- prediction = predictions[i]
- prediction = prediction.split()[-1]
- prediction = prediction[:-1] if self.enable_cot else prediction
- golden = goldens[i]
-
- # Define Metric
- score = self.scorer.exact_match_score(golden.expected_output, prediction)
- res.append({"prediction": prediction, "score": score})
-
- return res
-
-
-class BigBenchHardMetric:
- def __init__(
- self,
- model: GenerationMixin,
- tokenizer: PreTrainedTokenizerBase,
- device: str = "cuda",
- model_name: str = "Model",
- generation_params: dict = {},
- benchmark_params: dict = {},
- ) -> None:
- self.model = BigBenchHardModel(
- model, tokenizer, device, model_name, **generation_params
- )
- self.benchmark = BigBenchHardFixed(**benchmark_params)
-
- def score(self, batch_size: int | None = None) -> float:
- self.benchmark.evaluate(model=self.model, batch_size=batch_size)
- return self.benchmark.overall_score
diff --git a/src/atgen/metrics/compute_metrics.py b/src/atgen/metrics/compute_metrics.py
index 87e57f3..f55c175 100644
--- a/src/atgen/metrics/compute_metrics.py
+++ b/src/atgen/metrics/compute_metrics.py
@@ -1,5 +1,4 @@
from time import time
-from typing import List, Dict
import logging
from omegaconf import DictConfig
@@ -15,7 +14,7 @@
is_bart_score_available,
is_alignscore_available,
)
-from .supported_models_and_metrics import API_MODELS, DEEPEVAL_METRICS
+from .deepeval_supported_models_and_metrics import API_MODELS, DEEPEVAL_METRICS
log = logging.getLogger()
@@ -27,7 +26,7 @@ def compute_metrics(
original_texts,
config: DictConfig,
cache_dir: str = "cache",
-) -> Dict[str, float]:
+) -> dict[str, float]:
"""
Compute various metrics for generated texts.
@@ -90,7 +89,7 @@ def compute_metrics(
# Avoid division by zero
src_word_lengths_safe = np.where(src_word_lengths > 0, src_word_lengths, 1)
result["word_length_src_rel"] = result["word_length_gen"] / src_word_lengths_safe
- if "bartscore" in config.additional_metrics:
+ if "bartscore" in config.additional_metrics and is_bart_score_available:
log.info("Calculating BARTScore scores...")
start_time = time()
result.update(
@@ -138,24 +137,38 @@ def compute_metrics(
time_dict["time_rouge"] = time() - start_time
# Sacrebleu
start_time = time()
- sacrebleu_references = (
- [[ref] for ref in reference_texts]
- if not isinstance(reference_texts[0], list)
- else reference_texts
- )
- sacrebleu_result = sacrebleu.compute(
- predictions=generated_texts, references=sacrebleu_references
- )
- result["sacrebleu"] = sacrebleu_result.pop("score")
+ if not isinstance(reference_texts[0], list):
+ sacrebleu_references = [[ref] for ref in reference_texts]
+ sacrebleu_result = sacrebleu.compute(
+ predictions=generated_texts, references=sacrebleu_references
+ )
+ result["sacrebleu"] = sacrebleu_result.pop("score")
+ else:
+ sacrebleu_scores = []
+ for pred, ref in zip(generated_texts, reference_texts):
+ sacrebleu_result = sacrebleu.compute(
+ predictions=[pred], references=[ref]
+ )
+ sacrebleu_scores.append(sacrebleu_result.pop("score"))
+ result["sacrebleu"] = sacrebleu_scores
+
time_dict["time_sacrebleu"] = time() - start_time
# Lengths
- ref_word_lengths = np.array([len(text.split()) for text in reference_texts])
+ if isinstance(reference_texts[0], list):
+ ref_word_lengths = np.array(
+ [
+ np.mean([len(text.split()) for text in ref])
+ for ref in reference_texts
+ ]
+ )
+ else:
+ ref_word_lengths = np.array([len(ref.split()) for ref in reference_texts])
# Avoid division by zero
ref_word_lengths_safe = np.where(ref_word_lengths > 0, ref_word_lengths, 1)
result["word_length_rel"] = result["word_length_gen"] / ref_word_lengths_safe
# AlignScore
- if "alignscore" in config.additional_metrics:
+ if "alignscore" in config.additional_metrics and is_alignscore_available:
log.info("Calculating AlignScore scores...")
start_time = time()
alignscores = calculate_alignscore(
@@ -226,6 +239,5 @@ def compute_metrics(
and not "_reason" in key.lower()
and isinstance(value, (int, float)) # Ensure we only keep numerical metrics
}
- result.update(time_dict)
return result
diff --git a/src/atgen/metrics/supported_models_and_metrics.py b/src/atgen/metrics/deepeval_supported_models_and_metrics.py
similarity index 100%
rename from src/atgen/metrics/supported_models_and_metrics.py
rename to src/atgen/metrics/deepeval_supported_models_and_metrics.py
diff --git a/src/atgen/metrics/factory.py b/src/atgen/metrics/factory.py
new file mode 100644
index 0000000..1a42947
--- /dev/null
+++ b/src/atgen/metrics/factory.py
@@ -0,0 +1,270 @@
+from typing import Dict, List, Optional, Union, Any
+from dataclasses import dataclass, field
+import logging
+
+from .base import BaseMetric, MetricConfig
+from .lexical import BleuMetric, RougeMetric
+from .semantic import BartScoreMetric, AlignScoreMetric, SentBertMetric
+from .linguistic import ColaMetric
+from .llm_based import (
+ DeepEvalAnswerRelevancyMetric,
+ DeepEvalFaithfulnessMetric,
+ DeepEvalSummarizationMetric,
+ DeepEvalPromptAlignmentMetric,
+ BigBenchHardMetric,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MetricsConfig:
+
+ batch_size: int = 32
+ device: str = "cuda"
+ cache_dir: str = "cache"
+ aggregate: bool = True
+
+ metrics: List[str] = field(default_factory=list)
+
+ additional_metrics: List[str] = field(default_factory=list)
+
+ checkpoint: Optional[str] = None
+ model_name: Optional[str] = None
+
+ provider: Optional[str] = None
+ api_key: Optional[str] = None
+ base_url: Optional[str] = None
+ model: Optional[str] = None
+
+ threshold: float = 0.5
+ include_reason: bool = False
+ strict_mode: bool = False
+ async_mode: bool = True
+ verbose_mode: bool = False
+ truths_extraction_limit: Optional[int] = None
+
+ deepeval_threshold: Optional[float] = None
+ deepeval_include_reason: Optional[bool] = None
+ deepeval_strict_mode: Optional[bool] = None
+ deepeval_async_mode: Optional[bool] = None
+ deepeval_verbose_mode: Optional[bool] = None
+ deepeval_truths_extraction_limit: Optional[int] = None
+
+ benchmark_params: Dict[str, Any] = field(default_factory=dict)
+ generation_params: Dict[str, Any] = field(default_factory=dict)
+
+ custom_params: Dict[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self):
+ """Post-initialization to handle legacy parameter mappings."""
+ # Merge additional_metrics into metrics for backward compatibility
+ if self.additional_metrics:
+ self.metrics.extend(self.additional_metrics)
+ # Remove duplicates while preserving order
+ seen = set()
+ self.metrics = [x for x in self.metrics if not (x in seen or seen.add(x))]
+
+ # Handle DeepEval legacy parameter mappings
+ if self.deepeval_threshold is not None:
+ self.threshold = self.deepeval_threshold
+ if self.deepeval_include_reason is not None:
+ self.include_reason = self.deepeval_include_reason
+ if self.deepeval_strict_mode is not None:
+ self.strict_mode = self.deepeval_strict_mode
+ if self.deepeval_async_mode is not None:
+ self.async_mode = self.deepeval_async_mode
+ if self.deepeval_verbose_mode is not None:
+ self.verbose_mode = self.deepeval_verbose_mode
+ if self.deepeval_truths_extraction_limit is not None:
+ self.truths_extraction_limit = self.deepeval_truths_extraction_limit
+
+
+class MetricsFactory:
+ """Factory for creating metric instances."""
+
+ # Build registry dynamically to handle None values from import failures
+ _metric_registry = {}
+
+ @classmethod
+ def _build_registry(cls):
+ """Build the metric registry, excluding None values from failed imports."""
+ if not cls._metric_registry: # Only build once
+ registry = {
+ "bleu": BleuMetric,
+ "rouge1": RougeMetric,
+ "rouge2": RougeMetric,
+ "rougeL": RougeMetric,
+ "rougeLsum": RougeMetric,
+
+ "cola": ColaMetric,
+ "grammaticality": ColaMetric,
+
+ "deepeval_answer_relevance": DeepEvalAnswerRelevancyMetric,
+ "deepeval_faithfulness": DeepEvalFaithfulnessMetric,
+ "deepeval_summarization": DeepEvalSummarizationMetric,
+ "deepeval_prompt_alignment": DeepEvalPromptAlignmentMetric,
+ "bigbench_hard": BigBenchHardMetric,
+ "big_bench_hard": BigBenchHardMetric,
+ }
+
+ # Add semantic metrics if they imported successfully
+ if BartScoreMetric is not None:
+ registry["bartscore"] = BartScoreMetric
+ if AlignScoreMetric is not None:
+ registry["alignscore"] = AlignScoreMetric
+ if SentBertMetric is not None:
+ registry["sentbert"] = SentBertMetric
+ registry["sentence_bert"] = SentBertMetric
+
+ cls._metric_registry = registry
+
+ @classmethod
+ def register_metric(cls, name: str, metric_class: type):
+ """Register a new metric class."""
+ cls._build_registry()
+ cls._metric_registry[name] = metric_class
+ logger.info(f"Registered metric: {name} -> {metric_class.__name__}")
+
+ @classmethod
+ def get_available_metrics(cls) -> List[str]:
+ """Get list of all available metric names."""
+ cls._build_registry()
+ return list(cls._metric_registry.keys())
+
+ @classmethod
+ def get_all_possible_metric_keys(cls) -> List[str]:
+ """
+ Get list of all possible metric keys that can be returned by compute_metrics.
+ This includes not just metric names, but all the specific keys that metrics can return.
+ """
+ cls._build_registry()
+
+ # Base metric names
+ metric_keys = list(cls._metric_registry.keys())
+
+ # Add specific keys that metrics return (not just their names)
+ additional_keys = [
+ # ROUGE variants
+ "rouge1", "rouge2", "rougeL", "rougeLsum",
+
+ # SentBERT variants
+ "sentbert_pred_ref", "sentbert_pred_src",
+
+ # BLEU
+ "bleu",
+
+ # CoLA/Grammaticality
+ "cola", "grammaticality", "grammaticality_score",
+
+ # BARTScore variants (if available)
+ "bartscore", "bartscore_pred_ref", "bartscore_pred_src", "bartscore_ref_pred",
+
+ # AlignScore variants (if available)
+ "alignscore", "alignscore_pred_ref", "alignscore_pred_src",
+
+ # DeepEval metrics
+ "deepeval_answer_relevance", "deepeval_faithfulness",
+ "deepeval_summarization", "deepeval_prompt_alignment",
+
+ # BigBench variants
+ "bigbench_hard", "big_bench_hard",
+
+ # Statistical metrics that compute_metrics always adds
+ "word_length_gen", "word_length_src_rel", "word_length_rel", "exact_match",
+ ]
+
+ # Combine and deduplicate
+ all_keys = list(set(metric_keys + additional_keys))
+ all_keys.sort()
+
+ return all_keys
+
+ @classmethod
+ def create_metric(cls, name: str, config: Optional[MetricConfig] = None) -> BaseMetric:
+ cls._build_registry()
+ name = name.lower().strip()
+
+ if name not in cls._metric_registry:
+ available = ", ".join(cls.get_available_metrics())
+ raise ValueError(f"Unknown metric '{name}'. Available metrics: {available}")
+
+ metric_class = cls._metric_registry[name]
+ return metric_class(config)
+
+ @classmethod
+ def create_metrics(cls, metrics_config: MetricsConfig) -> Dict[str, BaseMetric]:
+ base_config = MetricConfig(
+ batch_size=metrics_config.batch_size,
+ device=metrics_config.device,
+ cache_dir=metrics_config.cache_dir,
+ aggregate=metrics_config.aggregate,
+ checkpoint=metrics_config.checkpoint,
+ model_name=metrics_config.model_name,
+ provider=metrics_config.provider,
+ api_key=metrics_config.api_key,
+ base_url=metrics_config.base_url,
+ model=metrics_config.model,
+ threshold=metrics_config.threshold,
+ include_reason=metrics_config.include_reason,
+ strict_mode=metrics_config.strict_mode,
+ async_mode=metrics_config.async_mode,
+ verbose_mode=metrics_config.verbose_mode,
+ truths_extraction_limit=metrics_config.truths_extraction_limit,
+ benchmark_params=metrics_config.benchmark_params,
+ generation_params=metrics_config.generation_params,
+ )
+
+ metrics = {}
+ for metric_name in metrics_config.metrics:
+ try:
+ metric = cls.create_metric(metric_name, base_config)
+ metrics[metric_name] = metric
+ logger.info(f"Created metric: {metric_name}")
+ except Exception as e:
+ logger.error(f"Failed to create metric '{metric_name}': {e}")
+
+ return metrics
+
+ @classmethod
+ def create_from_config(cls, config: Union[Dict[str, Any], MetricsConfig]) -> Dict[str, BaseMetric]:
+
+ if isinstance(config, dict):
+ metrics_config = MetricsConfig(**config)
+ else:
+ metrics_config = config
+
+ return cls.create_metrics(metrics_config)
+
+
+def get_metric_categories() -> Dict[str, List[str]]:
+ return {
+ "lexical": ["bleu", "rouge"],
+ "semantic": ["bartscore", "alignscore", "sentbert"],
+ "linguistic": ["cola"],
+ "llm_based": [
+ "deepeval_answer_relevance",
+ "deepeval_faithfulness",
+ "deepeval_summarization",
+ "deepeval_prompt_alignment"
+ ],
+ "benchmark": ["bigbench_hard"]
+ }
+
+
+def get_metric_requirements() -> Dict[str, Dict[str, bool]]:
+
+ return {
+ "bleu": {"requires_references": True, "requires_original_texts": False},
+ "rouge": {"requires_references": True, "requires_original_texts": False},
+ "bartscore": {"requires_references": True, "requires_original_texts": True},
+ "alignscore": {"requires_references": True, "requires_original_texts": True},
+ "sentbert": {"requires_references": False, "requires_original_texts": False},
+ "cola": {"requires_references": False, "requires_original_texts": False},
+ "deepeval_answer_relevance": {"requires_references": False, "requires_original_texts": True},
+ "deepeval_faithfulness": {"requires_references": True, "requires_original_texts": True},
+ "deepeval_summarization": {"requires_references": True, "requires_original_texts": True},
+ "deepeval_prompt_alignment": {"requires_references": False, "requires_original_texts": True},
+ "bigbench_hard": {"requires_references": False, "requires_original_texts": False},
+ "big_bench_hard": {"requires_references": False, "requires_original_texts": False},
+ }
\ No newline at end of file
diff --git a/src/atgen/metrics/lexical/__init__.py b/src/atgen/metrics/lexical/__init__.py
new file mode 100644
index 0000000..cab0fff
--- /dev/null
+++ b/src/atgen/metrics/lexical/__init__.py
@@ -0,0 +1,9 @@
+"""Lexical metrics for text evaluation."""
+
+from .bleu_metric import BleuMetric
+from .rouge_metric import RougeMetric
+
+__all__ = [
+ "BleuMetric",
+ "RougeMetric",
+]
\ No newline at end of file
diff --git a/src/atgen/metrics/lexical/bleu_metric.py b/src/atgen/metrics/lexical/bleu_metric.py
new file mode 100644
index 0000000..4440ccf
--- /dev/null
+++ b/src/atgen/metrics/lexical/bleu_metric.py
@@ -0,0 +1,110 @@
+from typing import Dict, List, Optional, Union
+import numpy as np
+from nltk import download
+from nltk.tokenize import word_tokenize
+from nltk.translate.bleu_score import corpus_bleu
+import logging
+
+from ..base.base_metric import BaseMetric, MetricConfig
+
+<<<<<<< HEAD
+=======
+# Ensure nltk data is downloaded
+>>>>>>> a24e0f2 (removed dependencies chek)
+try:
+ word_tokenize("test")
+except LookupError:
+ download('punkt')
+
+
+def smoothing_function(p_n, references, hypothesis, hyp_len):
+ """
+ Smooth-BLEU (BLEUS) as proposed in the paper:
+ Chin-Yew Lin, Franz Josef Och. ORANGE: a method for evaluating automatic
+ evaluation metrics for machine translation. COLING 2004.
+ """
+ smoothed_p_n = []
+ for i, p_i in enumerate(p_n, start=1):
+ # Smoothing is not applied for unigrams
+ if i > 1:
+ # If hypothesis length is lower than the current order, its value equals (0 + 1) / (0 + 1) = 0
+ if hyp_len < i:
+ assert p_i.denominator == 1
+ smoothed_p_n.append(1)
+ # Otherwise apply smoothing
+ else:
+ smoothed_p_i = (p_i.numerator + 1) / (p_i.denominator + 1)
+ smoothed_p_n.append(smoothed_p_i)
+ else:
+ smoothed_p_n.append(p_i)
+ return smoothed_p_n
+
+
+class BleuMetric(BaseMetric):
+<<<<<<< HEAD
+=======
+ """BLEU (Bilingual Evaluation Understudy) metric implementation."""
+>>>>>>> a24e0f2 (removed dependencies chek)
+
+ @property
+ def category(self) -> str:
+ return "lexical"
+
+ @property
+ def supports_multiple_references(self) -> bool:
+ return True
+
+
+<<<<<<< HEAD
+ def calculate(
+=======
+ def _compute(
+>>>>>>> a24e0f2 (removed dependencies chek)
+ self,
+ predictions: List[str],
+ references: Optional[List[Union[str, List[str]]]] = None,
+ original_texts: Optional[List[str]] = None
+ ) -> Dict[str, Union[float, np.ndarray]]:
+ """
+ Compute BLEU scores.
+
+ Args:
+ predictions: List of predicted texts
+ references: List of reference texts (can be list of lists for multiple references)
+ original_texts: Not used for BLEU computation
+
+ Returns:
+ Dictionary with BLEU scores
+ """
+ scores = []
+
+ for pred, ref in zip(predictions, references):
+ # Ensure references is always a list of lists
+ if isinstance(ref, str):
+ ref = [ref]
+
+ # Tokenize
+ tok_ref = [word_tokenize(r) for r in ref]
+ tok_pred = word_tokenize(pred)
+
+ try:
+<<<<<<< HEAD
+ score = corpus_bleu([tok_ref], [tok_pred], smoothing_function=smoothing_function)
+=======
+ score = corpus_bleu(tok_ref, [tok_pred], smoothing_function=smoothing_function)
+>>>>>>> a24e0f2 (removed dependencies chek)
+ scores.append(score)
+ except (KeyError, ZeroDivisionError):
+ scores.append(0.0)
+
+<<<<<<< HEAD
+ scores_dict = {"bleu": np.array(scores)}
+
+ # Apply aggregation if requested
+ if self.config.aggregate:
+ scores_dict = {key: float(np.mean(value)) for key, value in scores_dict.items()}
+
+ return scores_dict
+=======
+ return {"bleu": np.array(scores)}
+>>>>>>> a24e0f2 (removed dependencies chek)
diff --git a/src/atgen/metrics/lexical/rouge_metric.py b/src/atgen/metrics/lexical/rouge_metric.py
new file mode 100644
index 0000000..661dd08
--- /dev/null
+++ b/src/atgen/metrics/lexical/rouge_metric.py
@@ -0,0 +1,78 @@
+from typing import List, Dict, Union, Optional
+import numpy as np
+from evaluate import load
+
+from ..base.base_metric import BaseMetric, MetricConfig
+
+
+class RougeMetric(BaseMetric):
+ """ROUGE (Recall-Oriented Understudy for Gisting Evaluation) metric."""
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.rouge = None
+
+ def _initialize_rouge(self):
+ """Initialize ROUGE if not already initialized."""
+ if self.rouge is None:
+ self.rouge = load("rouge", cache_dir=self.config.cache_dir)
+<<<<<<< HEAD
+
+=======
+>>>>>>> a24e0f2 (removed dependencies chek)
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate ROUGE scores.
+
+ Args:
+ predictions: List of predicted texts
+ references: List of reference texts (can be list of lists for multiple references)
+ original_texts: Not used for ROUGE computation
+
+ Returns:
+ Dictionary with ROUGE scores
+ """
+
+ if references is None:
+ raise ValueError("ROUGE requires reference texts")
+
+ self._initialize_rouge()
+
+ # ROUGE expects different format for multiple references
+ if isinstance(references[0], list):
+ # Multiple references - convert to format expected by ROUGE
+ rouge_references = references
+ else:
+ # Single references
+ rouge_references = references
+
+ results = self.rouge.compute(
+ predictions=predictions,
+ references=rouge_references,
+ use_stemmer=True,
+ )
+
+ # Convert to float values
+ scores = {}
+ for key, value in results.items():
+ if isinstance(value, (int, float)):
+ scores[key] = float(value)
+<<<<<<< HEAD
+ elif hasattr(value, 'item'): #
+=======
+ elif hasattr(value, 'item'): # For numpy scalars
+>>>>>>> a24e0f2 (removed dependencies chek)
+ scores[key] = float(value.item())
+ else:
+ scores[key] = float(value)
+
+<<<<<<< HEAD
+ if self.config.aggregate:
+ for key, value in scores.items():
+ if isinstance(value, np.ndarray):
+ scores[key] = float(np.mean(value))
+
+=======
+>>>>>>> a24e0f2 (removed dependencies chek)
+ return scores
\ No newline at end of file
diff --git a/src/atgen/metrics/linguistic/__init__.py b/src/atgen/metrics/linguistic/__init__.py
new file mode 100644
index 0000000..b119e6d
--- /dev/null
+++ b/src/atgen/metrics/linguistic/__init__.py
@@ -0,0 +1,7 @@
+"""Linguistic quality metrics for text evaluation."""
+
+from .cola_metric import ColaMetric
+
+__all__ = [
+ "ColaMetric",
+]
\ No newline at end of file
diff --git a/src/atgen/metrics/linguistic/cola_metric.py b/src/atgen/metrics/linguistic/cola_metric.py
new file mode 100644
index 0000000..ec7e08a
--- /dev/null
+++ b/src/atgen/metrics/linguistic/cola_metric.py
@@ -0,0 +1,113 @@
+from typing import List, Dict, Union, Optional
+import numpy as np
+import torch
+from datasets import Dataset
+from torch.utils.data import DataLoader
+from transformers import (
+ AutoModelForSequenceClassification,
+ AutoTokenizer,
+ DataCollatorWithPadding,
+)
+import nltk
+
+from ..base.base_metric import BaseMetric, MetricConfig
+
+<<<<<<< HEAD
+=======
+# Ensure NLTK data is downloaded
+>>>>>>> a24e0f2 (removed dependencies chek)
+try:
+ nltk.data.find('tokenizers/punkt')
+except LookupError:
+ nltk.download('punkt')
+
+
+class ColaMetric(BaseMetric):
+ """CoLA (Corpus of Linguistic Acceptability) metric for evaluating grammaticality."""
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.model = None
+ self.tokenizer = None
+<<<<<<< HEAD
+ self.checkpoint = (config.checkpoint if config and config.checkpoint else 'Aktsvigun/electra-large-cola')
+=======
+ self.checkpoint = getattr(config, 'checkpoint', 'Aktsvigun/electra-large-cola') if config else 'Aktsvigun/electra-large-cola'
+>>>>>>> a24e0f2 (removed dependencies chek)
+
+
+ def _initialize_model(self):
+ """Initialize the CoLA model if not already initialized."""
+ if self.model is None:
+ self.model = AutoModelForSequenceClassification.from_pretrained(
+ self.checkpoint,
+ cache_dir=self.config.cache_dir
+ ).to(self.config.device)
+ self.tokenizer = AutoTokenizer.from_pretrained(
+ self.checkpoint,
+ cache_dir=self.config.cache_dir
+ )
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate CoLA grammaticality scores.
+
+ Args:
+ predictions: List of predicted texts to evaluate for grammaticality
+ references: Not used for CoLA
+ original_texts: Not used for CoLA
+
+ Returns:
+ Dictionary with CoLA scores
+ """
+
+ self._initialize_model()
+
+ # Split texts into sentences
+ text_sentences = [nltk.sent_tokenize(text) for text in predictions]
+ len_maps = np.cumsum([len(x) for x in text_sentences])
+ sentences = [sent for text in text_sentences for sent in text]
+
+ # Tokenize sentences
+ def tokenize_fn(instance):
+ return self.tokenizer(instance["text"], truncation=True)
+
+ tokenized_data = Dataset.from_dict({"text": sentences}).map(
+ tokenize_fn, remove_columns=["text"], batched=True
+ )
+ data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
+ dataloader = DataLoader(
+ tokenized_data,
+ batch_size=self.config.batch_size,
+ shuffle=False,
+ collate_fn=data_collator
+ )
+
+ # Calculate probabilities
+ sent_probas = torch.empty(len(sentences), dtype=torch.float32, device=self.config.device)
+ probas = torch.empty(len(predictions), dtype=torch.float32, device=self.config.device)
+
+ start = 0
+ end = self.config.batch_size
+
+ with torch.no_grad():
+ for i, batch in enumerate(dataloader):
+ batch = {k: v.to(self.config.device) for k, v in batch.items()}
+ batch_pred = self.model(**batch)
+ batch_probas = 1 / (1 + (-batch_pred.logits[:, 1]).exp())
+ sent_probas[start:end].copy_(batch_probas)
+ start = end
+ end += self.config.batch_size
+
+ # Aggregate sentence scores to text scores
+ for i, end_idx in enumerate(len_maps):
+ start_idx = len_maps[i - 1] if i != 0 else 0
+ probas[i].copy_(sent_probas[start_idx:end_idx].mean())
+
+ scores = {"cola": probas.cpu().detach().numpy()}
+
+ # Aggregate if requested
+ if self.config.aggregate:
+ scores = {key: float(np.mean(value)) for key, value in scores.items()}
+
+ return scores
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/__init__.py b/src/atgen/metrics/llm_based/__init__.py
new file mode 100644
index 0000000..45993a2
--- /dev/null
+++ b/src/atgen/metrics/llm_based/__init__.py
@@ -0,0 +1,19 @@
+"""LLM-based metrics for text evaluation using DeepEval."""
+
+from .evaluation_llm import EvaluationLLM
+from .base_deepeval_metric import BaseDeepEvalMetric
+from .answer_relevancy_metric import DeepEvalAnswerRelevancyMetric
+from .faithfulness_metric import DeepEvalFaithfulnessMetric
+from .summarization_metric import DeepEvalSummarizationMetric
+from .prompt_alignment_metric import DeepEvalPromptAlignmentMetric
+from .bigbench_hard_metric import BigBenchHardMetric
+
+__all__ = [
+ "EvaluationLLM",
+ "BaseDeepEvalMetric",
+ "DeepEvalAnswerRelevancyMetric",
+ "DeepEvalFaithfulnessMetric",
+ "DeepEvalSummarizationMetric",
+ "DeepEvalPromptAlignmentMetric",
+ "BigBenchHardMetric",
+]
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/answer_relevancy_metric.py b/src/atgen/metrics/llm_based/answer_relevancy_metric.py
new file mode 100644
index 0000000..c605375
--- /dev/null
+++ b/src/atgen/metrics/llm_based/answer_relevancy_metric.py
@@ -0,0 +1,46 @@
+"""
+DeepEval Answer Relevancy metric for evaluating how well outputs answer inputs.
+"""
+
+from typing import List, Dict, Union, Optional
+from deepeval.test_case import LLMTestCase
+from deepeval.metrics import AnswerRelevancyMetric
+
+from .base_deepeval_metric import BaseDeepEvalMetric
+
+
+class DeepEvalAnswerRelevancyMetric(BaseDeepEvalMetric):
+ """DeepEval Answer Relevancy metric for evaluating how well the output answers the input."""
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate Answer Relevancy scores.
+
+ Args:
+ predictions: List of predicted texts
+ references: Not used for Answer Relevancy
+ original_texts: List of original/source texts (required)
+
+ Returns:
+ Dictionary with Answer Relevancy scores
+ """
+ if not self.is_available():
+ raise RuntimeError("DeepEval dependencies not available")
+
+ if original_texts is None:
+ raise ValueError("Answer Relevancy metric requires original texts")
+
+ # Create DeepEval metric
+ metric = self._create_deepeval_metric(AnswerRelevancyMetric)
+
+ # Create test cases
+ test_cases = []
+ for pred, src in zip(predictions, original_texts):
+ test_case = LLMTestCase(
+ input=src,
+ actual_output=pred,
+ )
+ test_cases.append(test_case)
+
+ # Run evaluation
+ return self._run_evaluation(test_cases, metric, "deepeval_answer_relevance")
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/base_deepeval_metric.py b/src/atgen/metrics/llm_based/base_deepeval_metric.py
new file mode 100644
index 0000000..26b56b1
--- /dev/null
+++ b/src/atgen/metrics/llm_based/base_deepeval_metric.py
@@ -0,0 +1,131 @@
+"""
+Base class for DeepEval metrics.
+"""
+
+import os
+import sys
+from typing import List, Dict, Union, Optional
+import numpy as np
+from deepeval import evaluate
+from deepeval.test_case import LLMTestCase
+
+from ..base.base_metric import BaseMetric, MetricConfig
+from .evaluation_llm import EvaluationLLM
+
+
+class BaseDeepEvalMetric(BaseMetric):
+ """Base class for all DeepEval metrics."""
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.llm = None
+
+ def _validate_config(self):
+ """Validate DeepEval configuration."""
+ super()._validate_config()
+ if self.config.api_key is None:
+ self.logger.warning("API key is recommended for DeepEval metrics")
+
+ def _initialize_llm(self):
+ """Initialize the LLM for evaluation if not already initialized."""
+ if self.llm is None:
+<<<<<<< HEAD
+ # Determine base_url based on provider if not explicitly set
+ base_url = self.config.base_url
+ if base_url is None and self.config.provider:
+ provider = self.config.provider.lower()
+ if provider == "openai":
+ base_url = "https://api.openai.com/v1"
+ elif provider == "anthropic":
+ base_url = "https://api.anthropic.com/v1"
+ elif provider == "openrouter":
+ base_url = "https://openrouter.ai/api/v1"
+ else:
+ self.logger.warning(f"Unknown provider '{provider}', using default base_url")
+ base_url = "https://openrouter.ai/api/v1"
+ else:
+ base_url = base_url or "https://openrouter.ai/api/v1"
+
+ # Determine default model based on provider if not explicitly set
+ model = self.config.model
+ if model is None and self.config.provider:
+ provider = self.config.provider.lower()
+ if provider == "openai":
+ model = "gpt-4o-mini"
+ elif provider == "anthropic":
+ model = "claude-3-5-sonnet"
+ elif provider == "openrouter":
+ model = "openai/gpt-4o-mini"
+ else:
+ model = "openai/gpt-4o-mini"
+ else:
+ model = model or "openai/gpt-4o-mini"
+
+ self.llm = EvaluationLLM(
+ api_key=self.config.api_key,
+ model=model,
+ base_url=base_url,
+=======
+ self.llm = EvaluationLLM(
+ api_key=self.config.api_key,
+ model=self.config.model or "openai/gpt-4o-2024-11-20",
+ base_url=self.config.base_url or "https://openrouter.ai/api/v1",
+>>>>>>> a24e0f2 (removed dependencies chek)
+ )
+
+ def _create_deepeval_metric(self, metric_class, **kwargs):
+ """Create a DeepEval metric instance with common parameters."""
+ self._initialize_llm()
+
+ return metric_class(
+ threshold=self.config.threshold,
+ model=self.llm,
+ include_reason=self.config.include_reason,
+ strict_mode=self.config.strict_mode,
+ async_mode=self.config.async_mode,
+ **kwargs
+ )
+
+ def _run_evaluation(self, test_cases: List[LLMTestCase], metric, metric_name: str) -> Dict[str, float]:
+ """Run evaluation with proper output handling."""
+ if not test_cases:
+ return {}
+
+<<<<<<< HEAD
+=======
+ # Disable printing to console during evaluation if not verbose
+>>>>>>> a24e0f2 (removed dependencies chek)
+ original_stdout = sys.stdout
+ if not self.config.verbose_mode:
+ sys.stdout = open(os.devnull, "w")
+
+ try:
+<<<<<<< HEAD
+=======
+ # Run evaluation
+>>>>>>> a24e0f2 (removed dependencies chek)
+ evaluation_results = evaluate(
+ test_cases=test_cases,
+ metrics=[metric],
+ run_async=self.config.async_mode,
+ )
+
+ # Process results
+ scores = []
+ for result in evaluation_results.test_results:
+ scores.append(1 if result.success else 0)
+
+ # Calculate average score
+ if scores:
+ if self.config.aggregate:
+ return {metric_name: float(np.mean(scores))}
+ else:
+ return {metric_name: np.array(scores)}
+ else:
+ return {metric_name: 0.0}
+
+ finally:
+ # Restore stdout
+ if not self.config.verbose_mode:
+ sys.stdout.close()
+ sys.stdout = original_stdout
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/bigbench_hard_metric.py b/src/atgen/metrics/llm_based/bigbench_hard_metric.py
new file mode 100644
index 0000000..72dad02
--- /dev/null
+++ b/src/atgen/metrics/llm_based/bigbench_hard_metric.py
@@ -0,0 +1,186 @@
+"""
+BigBenchHard metric for evaluating model performance on challenging reasoning tasks.
+"""
+
+from typing import List, Dict, Union, Optional
+from transformers import GenerationMixin, PreTrainedTokenizerBase
+from deepeval.benchmarks import BigBenchHard
+from deepeval.benchmarks.big_bench_hard.template import BigBenchHardTemplate
+from deepeval.models.base_model import DeepEvalBaseLLM
+
+from ..base.base_metric import BaseMetric, MetricConfig
+
+
+class BigBenchHardModel(DeepEvalBaseLLM):
+ """DeepEval model wrapper for BigBenchHard evaluation."""
+
+ def __init__(
+ self,
+ model: GenerationMixin,
+ tokenizer: PreTrainedTokenizerBase,
+ device: str,
+ model_name: str = "Model",
+ **generation_kwargs,
+ ) -> None:
+ self.model = model
+ self.tokenizer = tokenizer
+ self.device = device
+ self.model_name = model_name
+ self.generation_kwargs = generation_kwargs
+
+ def load_model(self) -> GenerationMixin:
+ return self.model
+
+ def generate(self, prompt: str, **kwargs) -> str:
+ model = self.load_model()
+ model.to(self.device)
+ model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device)
+ generated_ids = model.generate(**model_inputs, **self.generation_kwargs)
+ return self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+ async def a_generate(self, prompt: str, **kwargs) -> str:
+ return self.generate(prompt)
+
+ def batch_generate(self, prompts: list[str], **kwargs) -> list[str]:
+ model = self.load_model()
+ model.to(self.device)
+ model_inputs = self.tokenizer(prompts, return_tensors="pt").to(self.device)
+ generated_ids = model.generate(**model_inputs, **self.generation_kwargs)
+ return self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+ def get_model_name(self):
+ return self.model_name
+
+
+class BigBenchHardFixed(BigBenchHard):
+ """
+ Fixed version of BigBenchHard that addresses batch prediction issues.
+
+ The main branch for deepeval contains broken code, as "NumberModel" is not defined:
+ https://github.com/confident-ai/deepeval/blob/main/deepeval/benchmarks/big_bench_hard/big_bench_hard.py#L153
+ """
+
+ def batch_predict(self, model, task, goldens):
+ prompts = []
+ for golden in goldens:
+ prompt: dict = BigBenchHardTemplate.generate_output(
+ input=golden.input,
+ task=task,
+ n_shots=self.n_shots,
+ enable_cot=self.enable_cot,
+ )
+ prompts.append(prompt)
+
+ # Enforced model generation
+ prompts = [
+ prompt + "Make sure to output only the numerical answer."
+ for prompt in prompts
+ ]
+ predictions = model.batch_generate(prompts)
+ predictions = [str(pred) for pred in predictions]
+
+ if len(predictions) != len(goldens):
+ raise ValueError(
+ "Custom `batch_generate` method did not return the same "
+ "number of generations as the number of prompts."
+ )
+
+ res = []
+ for i in range(len(predictions)):
+ prediction = predictions[i]
+ prediction = prediction.split()[-1]
+ prediction = prediction[:-1] if self.enable_cot else prediction
+ golden = goldens[i]
+
+ # Define Metric
+ score = self.scorer.exact_match_score(golden.expected_output, prediction)
+ res.append({"prediction": prediction, "score": score})
+
+ return res
+
+
+class BigBenchHardMetric(BaseMetric):
+ """BigBenchHard benchmark metric for evaluating reasoning performance."""
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.model = None
+ self.tokenizer = None
+ self.benchmark = None
+ self.benchmark_model = None
+
+ # Generation parameters
+ self.generation_params = getattr(config, 'generation_params', {}) if config else {}
+ self.benchmark_params = getattr(config, 'benchmark_params', {}) if config else {}
+ self.model_name = getattr(config, 'model_name', 'Model') if config else 'Model'
+
+
+ def set_model_and_tokenizer(self, model: GenerationMixin, tokenizer: PreTrainedTokenizerBase):
+ """
+ Set the model and tokenizer for evaluation.
+
+ Args:
+ model: The model to evaluate
+ tokenizer: The tokenizer for the model
+ """
+ self.model = model
+ self.tokenizer = tokenizer
+
+ # Initialize the benchmark model and benchmark
+ self.benchmark_model = BigBenchHardModel(
+ model=self.model,
+ tokenizer=self.tokenizer,
+ device=self.config.device,
+ model_name=self.model_name,
+ **self.generation_params
+ )
+ self.benchmark = BigBenchHardFixed(**self.benchmark_params)
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate BigBenchHard score.
+
+ Note: This method signature is maintained for consistency with BaseMetric,
+ but BigBenchHard doesn't use predictions/references in the traditional way.
+ The model and tokenizer must be set using set_model_and_tokenizer() before calling this.
+
+ Args:
+ predictions: Not used for BigBenchHard
+ references: Not used for BigBenchHard
+ original_texts: Not used for BigBenchHard
+
+ Returns:
+ Dictionary with BigBenchHard score
+ """
+ if not self.is_available():
+ raise RuntimeError("BigBenchHard dependencies not available")
+
+ if self.model is None or self.tokenizer is None:
+ raise ValueError("Model and tokenizer must be set using set_model_and_tokenizer() before evaluation")
+
+ if self.benchmark_model is None or self.benchmark is None:
+ self.set_model_and_tokenizer(self.model, self.tokenizer)
+
+ # Run the benchmark evaluation
+ self.benchmark.evaluate(model=self.benchmark_model, batch_size=self.config.batch_size)
+
+ return {"bigbench_hard_score": float(self.benchmark.overall_score)}
+
+ def evaluate_benchmark(self, model: GenerationMixin, tokenizer: PreTrainedTokenizerBase, batch_size: Optional[int] = None) -> float:
+ """
+ Convenience method for direct benchmark evaluation.
+
+ Args:
+ model: The model to evaluate
+ tokenizer: The tokenizer for the model
+ batch_size: Batch size for evaluation (uses config default if None)
+
+ Returns:
+ Overall benchmark score
+ """
+ self.set_model_and_tokenizer(model, tokenizer)
+
+ eval_batch_size = batch_size if batch_size is not None else self.config.batch_size
+ self.benchmark.evaluate(model=self.benchmark_model, batch_size=eval_batch_size)
+
+ return self.benchmark.overall_score
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/evaluation_llm.py b/src/atgen/metrics/llm_based/evaluation_llm.py
new file mode 100644
index 0000000..556ca79
--- /dev/null
+++ b/src/atgen/metrics/llm_based/evaluation_llm.py
@@ -0,0 +1,94 @@
+"""
+Custom Evaluation LLM implementation for DeepEval metrics.
+"""
+
+from openai import OpenAI, AsyncOpenAI
+from deepeval.models.base_model import DeepEvalBaseLLM
+
+
+class EvaluationLLM(DeepEvalBaseLLM):
+ """
+ Custom Evaluation LLM implementation for DeepEval.
+
+ This class implements the DeepEvalBaseLLM interface to allow using
+ custom models with DeepEval metrics.
+ """
+
+ def __init__(
+ self,
+ api_key=None,
+ model="openai/gpt-4o-2024-11-20",
+ base_url="https://openrouter.ai/api/v1",
+ ):
+ """
+ Initialize the Evaluation LLM.
+
+ Args:
+ api_key: Evaluation API key
+ model: Model identifier (e.g., "openai/gpt-4o-2024-11-20")
+ base_url: Evaluation API base URL
+ """
+ self.api_key = api_key
+ self.model_name = model
+ self.base_url = base_url
+ self.client = None
+ self.async_client = None
+ self.OpenAI = OpenAI
+ self.AsyncOpenAI = AsyncOpenAI
+
+ def load_model(self):
+ """Load and return the client."""
+ if self.client is None:
+ self.client = self.OpenAI(
+ base_url=self.base_url,
+ api_key=self.api_key,
+ )
+ return self.client
+
+ def load_async_model(self):
+ """Load and return the async client."""
+ if self.async_client is None:
+ self.async_client = self.AsyncOpenAI(
+ base_url=self.base_url,
+ api_key=self.api_key,
+ )
+ return self.async_client
+
+ def generate(self, prompt: str) -> str:
+ """
+ Generate a response from the evaluation model.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+ """
+ client = self.load_model()
+ response = client.chat.completions.create(
+ model=self.model_name,
+ messages=[{"role": "user", "content": prompt}],
+ )
+ return response.choices[0].message.content
+
+ async def a_generate(self, prompt: str) -> str:
+ """
+ Asynchronously generate a response from the evaluation model.
+
+ Args:
+ prompt: The prompt to send to the model
+
+ Returns:
+ The model's response as a string
+ """
+ # Use the async client for async operations
+ client = self.load_async_model()
+ response = await client.chat.completions.create(
+ model=self.model_name,
+ messages=[{"role": "user", "content": prompt}],
+ )
+ return response.choices[0].message.content
+
+ def get_model_name(self):
+ """Return the name of the model."""
+ return f"EvaluationLLM: {self.model_name}"
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/faithfulness_metric.py b/src/atgen/metrics/llm_based/faithfulness_metric.py
new file mode 100644
index 0000000..9f42e07
--- /dev/null
+++ b/src/atgen/metrics/llm_based/faithfulness_metric.py
@@ -0,0 +1,51 @@
+"""
+DeepEval Faithfulness metric for evaluating factual consistency with input.
+"""
+
+from typing import List, Dict, Union, Optional
+from deepeval.test_case import LLMTestCase
+from deepeval.metrics import FaithfulnessMetric
+
+from .base_deepeval_metric import BaseDeepEvalMetric
+
+
+class DeepEvalFaithfulnessMetric(BaseDeepEvalMetric):
+ """DeepEval Faithfulness metric for evaluating factual consistency with the input."""
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate Faithfulness scores.
+
+ Args:
+ predictions: List of predicted texts
+ references: Not used for Faithfulness
+ original_texts: List of original/source texts (required)
+
+ Returns:
+ Dictionary with Faithfulness scores
+ """
+ if not self.is_available():
+ raise RuntimeError("DeepEval dependencies not available")
+
+ if original_texts is None:
+ raise ValueError("Faithfulness metric requires original texts")
+
+ # Create DeepEval metric with truths extraction limit if specified
+ kwargs = {}
+ if self.config.truths_extraction_limit is not None:
+ kwargs['truths_extraction_limit'] = self.config.truths_extraction_limit
+
+ metric = self._create_deepeval_metric(FaithfulnessMetric, **kwargs)
+
+ # Create test cases
+ test_cases = []
+ for pred, src in zip(predictions, original_texts):
+ test_case = LLMTestCase(
+ input=src,
+ actual_output=pred,
+ retrieval_context=[src], # Use source as retrieval context
+ )
+ test_cases.append(test_case)
+
+ # Run evaluation
+ return self._run_evaluation(test_cases, metric, "deepeval_faithfulness")
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/prompt_alignment_metric.py b/src/atgen/metrics/llm_based/prompt_alignment_metric.py
new file mode 100644
index 0000000..e549d18
--- /dev/null
+++ b/src/atgen/metrics/llm_based/prompt_alignment_metric.py
@@ -0,0 +1,55 @@
+"""
+DeepEval Prompt Alignment metric for evaluating alignment with expected output.
+"""
+
+from typing import List, Dict, Union, Optional
+from deepeval.test_case import LLMTestCase
+from deepeval.metrics import PromptAlignmentMetric
+
+from .base_deepeval_metric import BaseDeepEvalMetric
+
+
+class DeepEvalPromptAlignmentMetric(BaseDeepEvalMetric):
+ """DeepEval Prompt Alignment metric for evaluating alignment with the expected output."""
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate Prompt Alignment scores.
+
+ Args:
+ predictions: List of predicted texts
+ references: List of reference texts (required)
+ original_texts: List of original/source texts (required)
+
+ Returns:
+ Dictionary with Prompt Alignment scores
+ """
+
+ if original_texts is None:
+ raise ValueError("Prompt Alignment metric requires original texts")
+
+ if references is None:
+ raise ValueError("Prompt Alignment metric requires reference texts")
+
+ if isinstance(references[0], list):
+ self.logger.error("Prompt Alignment does not support multiple references. Skipping...")
+ return {}
+
+ # Create DeepEval metric with default prompt instructions
+ metric = self._create_deepeval_metric(
+ PromptAlignmentMetric,
+ prompt_instructions=["Do what you are told to do in the prompt"]
+ )
+
+ # Create test cases
+ test_cases = []
+ for pred, ref, src in zip(predictions, references, original_texts):
+ test_case = LLMTestCase(
+ input=src,
+ actual_output=pred,
+ expected_output=ref,
+ )
+ test_cases.append(test_case)
+
+ # Run evaluation
+ return self._run_evaluation(test_cases, metric, "deepeval_prompt_alignment")
\ No newline at end of file
diff --git a/src/atgen/metrics/llm_based/summarization_metric.py b/src/atgen/metrics/llm_based/summarization_metric.py
new file mode 100644
index 0000000..e5b94c6
--- /dev/null
+++ b/src/atgen/metrics/llm_based/summarization_metric.py
@@ -0,0 +1,44 @@
+"""
+DeepEval Summarization metric for evaluating summarization quality.
+"""
+
+from typing import List, Dict, Union, Optional
+from deepeval.test_case import LLMTestCase
+from deepeval.metrics import SummarizationMetric
+
+from .base_deepeval_metric import BaseDeepEvalMetric
+
+
+class DeepEvalSummarizationMetric(BaseDeepEvalMetric):
+ """DeepEval Summarization metric for evaluating summarization quality."""
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate Summarization scores.
+
+ Args:
+ predictions: List of predicted texts
+ references: Not used for Summarization
+ original_texts: List of original/source texts (required)
+
+ Returns:
+ Dictionary with Summarization scores
+ """
+
+ if original_texts is None:
+ raise ValueError("Summarization metric requires original texts")
+
+ # Create DeepEval metric
+ metric = self._create_deepeval_metric(SummarizationMetric)
+
+ # Create test cases
+ test_cases = []
+ for pred, src in zip(predictions, original_texts):
+ test_case = LLMTestCase(
+ input=src,
+ actual_output=pred,
+ )
+ test_cases.append(test_case)
+
+ # Run evaluation
+ return self._run_evaluation(test_cases, metric, "deepeval_summarization")
\ No newline at end of file
diff --git a/src/atgen/metrics/metrics.py b/src/atgen/metrics/metrics.py
index 551b4f5..52f9504 100644
--- a/src/atgen/metrics/metrics.py
+++ b/src/atgen/metrics/metrics.py
@@ -2,7 +2,7 @@
import os
import sys
from openai import OpenAI, AsyncOpenAI
-from typing import List, Dict
+from typing import Union
from urllib.request import urlretrieve
import logging
from pathlib import Path
@@ -62,7 +62,7 @@
# Going up 3 levels from metrics.py: src/atgen/metrics -> repository root
os.path.join(
Path(__file__).parents[3],
- "external_metrics/AlignScore/model/AlignScore-base.ckpt",
+ "cache/AlignScore-base.ckpt",
),
)
@@ -111,19 +111,20 @@ def smoothing_function(p_n, references, hypothesis, hyp_len):
return smoothed_p_n
-def pair_bleu(references, prediction):
+def pair_bleu(references: list[str] | str, prediction: str):
"""
Compute the bleu score between two given texts.
A smoothing function is used to avoid zero scores when
there are no common higher order n-grams between the
texts.
"""
- if not isinstance(references, list):
- references = [references]
- tok_ref = [word_tokenize(ref) for ref in references]
- tok_pred = word_tokenize(prediction)
+ if isinstance(references, str):
+ tok_ref = [[word_tokenize(references)]]
+ else:
+ tok_ref = [[word_tokenize(ref) for ref in references]]
+ tok_pred = [word_tokenize(prediction)]
try:
- return corpus_bleu(tok_ref, [tok_pred], smoothing_function=smoothing_function)
+ return corpus_bleu(tok_ref, tok_pred, smoothing_function=smoothing_function)
except (KeyError, ZeroDivisionError):
return 0.0
@@ -167,107 +168,6 @@ def calculate_bart_score(
return scores
-def calculate_summac_score(
- predictions: List[str],
- texts: List[str],
- labels: List[str] = None,
- aggregate: bool = True,
-) -> Dict[str, np.ndarray]:
- scorer = SummaCZS(granularity="sentence", model_name="vitc")
- preds_score = scorer.score(texts, predictions)["scores"]
- if labels is not None:
- labels_score = scorer.score(texts, labels)["scores"]
- rel_score = np.array(preds_score) / np.array(labels_score)
- if aggregate:
- preds_score = np.mean(preds_score)
- if labels is not None:
- rel_score = np.mean(rel_score)
- if labels is not None:
- return {"SummaC-tp": preds_score, "SummaC-rel": rel_score}
- return {"SummaC-tp": preds_score}
-
-
-def calculate_cola_model_predictions(
- texts,
- checkpoint="Aktsvigun/electra-large-cola",
- batch_size=64,
- device="cuda",
- return_sent_data: bool = False,
- aggregate: bool = True,
- cache_dir: str = "cache",
-):
- model = AutoModelForSequenceClassification.from_pretrained(
- "Aktsvigun/electra-large-cola"
- ).to(device)
- tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-
- text_sentences = [nltk.sent_tokenize(text) for text in texts]
- len_maps = np.cumsum([len(x) for x in text_sentences])
- sentences = [sent for text in text_sentences for sent in text]
-
- def tokenize_fn(instance):
- return tokenizer(instance["text"], truncation=True)
-
- tokenized_data = Dataset.from_dict({"text": sentences}).map(
- tokenize_fn, remove_columns=["text"], batched=True
- )
- data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
- dataloader = DataLoader(
- tokenized_data, batch_size=batch_size, shuffle=False, collate_fn=data_collator
- )
-
- sent_probas = torch.empty(len(sentences), dtype=torch.float32, device=device)
- probas = torch.empty(len(texts), dtype=torch.float32, device=device)
- start = 0
- end = batch_size
- with torch.no_grad():
- for i, batch in enumerate(dataloader):
- batch_pred = model(**{k: v.cuda() for k, v in batch.items()})
- batch_probas = 1 / (1 + (-batch_pred.logits[:, 1]).exp())
- sent_probas[start:end].copy_(batch_probas)
- start = end
- end += batch_size
-
- for i, end_idx in enumerate(len_maps):
- start_idx = len_maps[i - 1] if i != 0 else 0
- probas[i].copy_(sent_probas[start_idx:end_idx].mean())
-
- if aggregate:
- return probas.mean().item()
- if return_sent_data:
- return (
- probas.cpu().detach().numpy(),
- sent_probas.cpu().detach().numpy(),
- sentences,
- )
- return probas.cpu().detach().numpy()
-
-
-def calculate_infolm_score(predictions, references, batch_size=4):
- from torchmetrics.text.infolm import InfoLM
-
- assert len(predictions) == len(references), "Lengths must coincide!"
- infolm_fr = InfoLM(measure_to_use="fisher_rao")
- infolm_ab = InfoLM(measure_to_use="ab", alpha=1.0, beta=1.0)
-
- infolm_fr_scores, infolm_ab_scores = [], []
- idf_ref, idf_hyps = infolm_fr.prepare_idfs(references, predictions)
-
- num_batches = ceil(len(predictions) / batch_size)
- for i in range(num_batches):
- batch_preds = predictions[i * batch_size : (i + 1) * batch_size]
- batch_refs = references[i * batch_size : (i + 1) * batch_size]
-
- infolm_fr_scores += infolm_fr.evaluate_batch(
- batch_preds, batch_refs, idf_ref=idf_ref, idf_hyps=idf_hyps
- )["fisher_rao"]
- infolm_ab_scores = infolm_ab.evaluate_batch(
- batch_preds, batch_refs, idf_ref=idf_ref, idf_hyps=idf_hyps
- )["ab"]
-
- return infolm_fr_scores, infolm_ab_scores
-
-
def calculate_abstractiveness_scores(
predictions, texts, references=None, aggregate: bool = True
):
@@ -350,7 +250,7 @@ def __init__(
self.device = device
def __call__(
- self, source_texts: List[str], ref_texts: List[str], batch_size: int = 32
+ self, source_texts: list[str], ref_texts: list[str], batch_size: int = 32
) -> np.ndarray:
assert len(source_texts) == len(ref_texts)
# Make batch_size an even number
@@ -410,9 +310,9 @@ def mean_pooling(model_output, attention_mask):
def calculate_alignscore(
- predictions,
- references,
- original_texts,
+ predictions: list[str],
+ references: Union[list[str], list[list[str]]],
+ original_texts: list[str],
batch_size: int = 32,
device: str = "cuda",
cache_dir: str = "cache",
@@ -441,7 +341,15 @@ def calculate_alignscore(
references = [text if text else " " for text in references]
scores_ref = scorer.score(contexts=original_texts, claims=predictions)
- scores_baseline = scorer.score(contexts=original_texts, claims=references)
+ if isinstance(references[0], list):
+ scores_baseline = []
+ for orig_text, refs in zip(original_texts, references):
+ inst_baseline_scores = scorer.score(
+ contexts=[orig_text] * len(refs), claims=refs
+ )
+ scores_baseline.append(max(inst_baseline_scores))
+ else:
+ scores_baseline = scorer.score(contexts=original_texts, claims=references)
scores_rel = np.array(scores_ref) / np.array(scores_baseline)
return {"alignscore": scores_ref, "alignscore_rel": scores_rel}
@@ -554,10 +462,10 @@ def calculate_deepeval_metrics(
Calculate DeepEval metrics using EvaluationLLM.
Args:
- predictions: List of generated texts
- references: List of reference texts
- original_texts: List of source texts
- metrics_to_calculate: List of metrics to calculate. Options:
+ predictions: list of generated texts
+ references: list of reference texts
+ original_texts: list of source texts
+ metrics_to_calculate: list of metrics to calculate. Options:
["deepeval_answer_relevance", "deepeval_faithfulness", "deepeval_summarization", "deepeval_prompt_alignment"]
api_key: Evaluation API key
base_url: Evaluation API base URL
@@ -570,7 +478,7 @@ def calculate_deepeval_metrics(
truths_extraction_limit: Maximum number of factual truths to extract (default: None)
Returns:
- Dictionary with metric scores
+ dictionary with metric scores
"""
if not metrics_to_calculate:
diff --git a/src/atgen/metrics/semantic/__init__.py b/src/atgen/metrics/semantic/__init__.py
new file mode 100644
index 0000000..1d8da41
--- /dev/null
+++ b/src/atgen/metrics/semantic/__init__.py
@@ -0,0 +1,27 @@
+"""Semantic metrics for text evaluation."""
+
+import warnings
+
+# Import with error handling for optional dependencies
+__all__ = []
+
+try:
+ from .bart_score_metric import BartScoreMetric
+ __all__.append("BartScoreMetric")
+except ImportError as e:
+ warnings.warn(f"BartScore metric not available due to import error: {e}")
+ BartScoreMetric = None
+
+try:
+ from .alignscore_metric import AlignScoreMetric
+ __all__.append("AlignScoreMetric")
+except ImportError as e:
+ warnings.warn(f"AlignScore metric not available due to import error: {e}")
+ AlignScoreMetric = None
+
+try:
+ from .sentbert_metric import SentBertMetric
+ __all__.append("SentBertMetric")
+except ImportError as e:
+ warnings.warn(f"SentBert metric not available due to import error: {e}")
+ SentBertMetric = None
\ No newline at end of file
diff --git a/src/atgen/metrics/semantic/alignscore_metric.py b/src/atgen/metrics/semantic/alignscore_metric.py
new file mode 100644
index 0000000..2fffda2
--- /dev/null
+++ b/src/atgen/metrics/semantic/alignscore_metric.py
@@ -0,0 +1,90 @@
+import os
+from typing import List, Dict, Union, Optional
+from urllib.request import urlretrieve
+from pathlib import Path
+import numpy as np
+from alignscore import AlignScore
+from ..base.base_metric import BaseMetric, MetricConfig
+
+# Default AlignScore checkpoint path
+ALIGNSCORE_CHECKPOINT_PATH = os.getenv(
+ "ALIGNSCORE_CHECKPOINT_PATH",
+ os.path.join(
+ Path(__file__).parents[4], # Going up to repository root
+ "external_metrics/AlignScore/model/AlignScore-base.ckpt",
+ ),
+)
+
+
+class AlignScoreMetric(BaseMetric):
+ """AlignScore metric for evaluating factual consistency."""
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.scorer = None
+ self.checkpoint_path = getattr(config, 'checkpoint_path', ALIGNSCORE_CHECKPOINT_PATH) if config else ALIGNSCORE_CHECKPOINT_PATH
+
+
+ def _initialize_scorer(self):
+ """Initialize the AlignScore scorer if not already initialized."""
+ if self.scorer is None:
+ if not os.path.exists(self.checkpoint_path):
+ # Download the checkpoint if it doesn't exist
+ os.makedirs(os.path.dirname(self.checkpoint_path), exist_ok=True)
+ urlretrieve(
+ "https://huggingface.co/yzha/AlignScore/resolve/main/AlignScore-base.ckpt",
+ self.checkpoint_path,
+ )
+
+ self.scorer = AlignScore(
+ model="roberta-base",
+ batch_size=self.config.batch_size,
+ device=self.config.device,
+ ckpt_path=self.checkpoint_path,
+ evaluation_mode="nli_sp",
+ )
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate AlignScore metrics.
+
+ Args:
+ predictions: List of predicted texts
+ references: List of reference texts (single references only)
+ original_texts: List of original/source texts (required)
+
+ Returns:
+ Dictionary with AlignScore results
+ """
+ if original_texts is None:
+ raise ValueError("AlignScore requires original texts")
+
+ if references is not None and isinstance(references[0], list):
+ self.logger.error("AlignScore does not support multiple references. Skipping...")
+ return {}
+
+ self._initialize_scorer()
+
+ # Fix: AlignScore outputs an error if a text is empty, so we need to add some content
+ original_texts = [text if text else " " for text in original_texts]
+ predictions = [text if text else " " for text in predictions]
+
+ scores = {}
+
+ # Calculate AlignScore between predictions and original texts
+ scores_pred = self.scorer.score(contexts=original_texts, claims=predictions)
+ scores["alignscore"] = scores_pred
+
+ if references is not None:
+ references = [text if text else " " for text in references]
+ # Calculate baseline AlignScore between references and original texts
+ scores_baseline = self.scorer.score(contexts=original_texts, claims=references)
+ # Calculate relative score
+ scores_rel = np.array(scores_pred) / np.array(scores_baseline)
+ scores["alignscore_rel"] = scores_rel
+
+ # Aggregate if requested
+ if self.config.aggregate:
+ scores = {key: float(np.mean(value)) for key, value in scores.items()}
+
+ return scores
\ No newline at end of file
diff --git a/src/atgen/metrics/semantic/bart_score_metric.py b/src/atgen/metrics/semantic/bart_score_metric.py
new file mode 100644
index 0000000..afea7a3
--- /dev/null
+++ b/src/atgen/metrics/semantic/bart_score_metric.py
@@ -0,0 +1,152 @@
+import traceback
+from typing import List, Dict, Union, Optional
+import numpy as np
+import torch
+import torch.nn as nn
+from transformers import BartTokenizer, BartForConditionalGeneration
+from tqdm import tqdm
+
+from ..base.base_metric import BaseMetric, MetricConfig
+
+
+class BARTScorer:
+ """Internal BARTScorer implementation."""
+
+ def __init__(
+ self,
+ device="cuda",
+ max_length=1024,
+ checkpoint="facebook/bart-large-cnn",
+ cache_dir: str = "cache",
+ ):
+ # Set up model
+ self.device = device
+ self.max_length = max_length
+ self.tokenizer = BartTokenizer.from_pretrained(checkpoint, cache_dir=cache_dir)
+ self.model = BartForConditionalGeneration.from_pretrained(
+ checkpoint, cache_dir=cache_dir
+ )
+ self.model.eval()
+ self.model.to(device)
+
+ # Set up loss
+ self.loss_fct = nn.NLLLoss(
+ reduction="none", ignore_index=self.model.config.pad_token_id
+ )
+ self.lsm = nn.LogSoftmax(dim=1)
+
+ def score(self, srcs, tgts, batch_size=4):
+ """Score a batch of examples"""
+ score_list = []
+<<<<<<< HEAD
+ for i in range(0, len(srcs), batch_size):
+=======
+ for i in tqdm(range(0, len(srcs), batch_size), desc="Calculating BARTScore..."):
+>>>>>>> a24e0f2 (removed dependencies chek)
+ src_list = srcs[i : i + batch_size]
+ tgt_list = tgts[i : i + batch_size]
+ try:
+ with torch.no_grad():
+ encoded_src = self.tokenizer(
+ src_list,
+ max_length=self.max_length,
+ truncation=True,
+ padding=True,
+ return_tensors="pt",
+ )
+ encoded_tgt = self.tokenizer(
+ tgt_list,
+ max_length=self.max_length,
+ truncation=True,
+ padding=True,
+ return_tensors="pt",
+ )
+ src_tokens = encoded_src["input_ids"].to(self.device)
+ src_mask = encoded_src["attention_mask"].to(self.device)
+
+ tgt_tokens = encoded_tgt["input_ids"].to(self.device)
+ tgt_mask = encoded_tgt["attention_mask"]
+ tgt_len = tgt_mask.sum(dim=1).to(self.device)
+
+ output = self.model(
+ input_ids=src_tokens, attention_mask=src_mask, labels=tgt_tokens
+ )
+ logits = output.logits.view(-1, self.model.config.vocab_size)
+ loss = self.loss_fct(self.lsm(logits), tgt_tokens.view(-1))
+ loss = loss.view(tgt_tokens.shape[0], -1)
+ loss = loss.sum(dim=1) / tgt_len
+ curr_score_list = [-x.item() for x in loss]
+ score_list += curr_score_list
+
+ except RuntimeError:
+ traceback.print_exc()
+ print(f"source: {src_list}")
+ print(f"target: {tgt_list}")
+ exit(0)
+ return score_list
+
+
+class BartScoreMetric(BaseMetric):
+ """BARTScore metric for evaluating text generation quality."""
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.scorer = None
+ self.checkpoint = getattr(config, 'checkpoint', 'facebook/bart-large-cnn') if config else 'facebook/bart-large-cnn'
+
+
+ def _initialize_scorer(self):
+ """Initialize the BARTScorer if not already initialized."""
+ if self.scorer is None:
+ self.scorer = BARTScorer(
+ device=self.config.device,
+ checkpoint=self.checkpoint,
+ cache_dir=self.config.cache_dir,
+ )
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate BARTScore metrics.
+
+ Args:
+ predictions: List of predicted texts
+ references: List of reference texts (can be list of lists for multiple references)
+ original_texts: List of original/source texts
+
+ Returns:
+ Dictionary with BARTScore results
+ """
+<<<<<<< HEAD
+=======
+ if not self.is_available():
+ raise RuntimeError("BARTScore dependencies not available")
+>>>>>>> a24e0f2 (removed dependencies chek)
+
+ self._initialize_scorer()
+
+ scores = {}
+
+ # Calculate source-hypothesis score (BARTScore-sh)
+ if original_texts is not None:
+ sh_scores = self.scorer.score(original_texts, predictions, batch_size=self.config.batch_size)
+ scores["BARTScore-sh"] = np.array(sh_scores)
+
+ # Calculate hypothesis-reference score (BARTScore-hr)
+ if references is not None:
+ if isinstance(references[0], list):
+ # Handle multiple references - take maximum score
+ hr_scores = []
+ for ref, pred in zip(references, predictions):
+ inst_pred = [pred for _ in range(len(ref))]
+ inst_scores = self.scorer.score(inst_pred, ref, batch_size=self.config.batch_size)
+ hr_scores.append(max(inst_scores))
+ scores["BARTScore-hr"] = np.array(hr_scores)
+ else:
+ hr_scores = self.scorer.score(predictions, references, batch_size=self.config.batch_size)
+ scores["BARTScore-hr"] = np.array(hr_scores)
+
+ # Aggregate if requested
+ if self.config.aggregate:
+ scores = {key: float(np.mean(value)) for key, value in scores.items()}
+
+ return scores
\ No newline at end of file
diff --git a/src/atgen/metrics/semantic/sentbert_metric.py b/src/atgen/metrics/semantic/sentbert_metric.py
new file mode 100644
index 0000000..dff5450
--- /dev/null
+++ b/src/atgen/metrics/semantic/sentbert_metric.py
@@ -0,0 +1,163 @@
+from typing import List, Dict, Union, Optional
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel
+
+from ..base.base_metric import BaseMetric, MetricConfig
+
+
+class SentBertMetric(BaseMetric):
+<<<<<<< HEAD
+=======
+ """SentenceBERT semantic similarity metric."""
+>>>>>>> a24e0f2 (removed dependencies chek)
+
+ def __init__(self, config: Optional[MetricConfig] = None):
+ super().__init__(config)
+ self.model = None
+ self.tokenizer = None
+<<<<<<< HEAD
+ self.checkpoint = (config.checkpoint if config and config.checkpoint else 'sentence-transformers/all-mpnet-base-v2')
+
+
+ def _initialize_model(self):
+=======
+ self.checkpoint = getattr(config, 'checkpoint', 'sentence-transformers/all-mpnet-base-v2') if config else 'sentence-transformers/all-mpnet-base-v2'
+
+
+ def _initialize_model(self):
+ """Initialize the SentenceBERT model if not already initialized."""
+>>>>>>> a24e0f2 (removed dependencies chek)
+ if self.model is None:
+ self.tokenizer = AutoTokenizer.from_pretrained(
+ self.checkpoint,
+ cache_dir=self.config.cache_dir
+ )
+ self.model = AutoModel.from_pretrained(
+ self.checkpoint,
+ cache_dir=self.config.cache_dir
+ ).to(self.config.device)
+
+ @staticmethod
+ def mean_pooling(model_output, attention_mask):
+ """Mean Pooling - Take attention mask into account for correct averaging."""
+ token_embeddings = model_output[0] # First element contains all token embeddings
+ input_mask_expanded = (
+ attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+ )
+ return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+ input_mask_expanded.sum(1), min=1e-9
+ )
+
+ def _compute_similarity(self, source_texts: List[str], ref_texts: List[str]) -> np.ndarray:
+ """Compute semantic similarity between source and reference texts."""
+ assert len(source_texts) == len(ref_texts)
+
+ # Make batch_size an odd number for better batch processing
+ batch_size = self.config.batch_size
+ if batch_size % 2 == 0:
+ batch_size -= 1
+ half_batch_size = batch_size // 2
+ n_texts = len(source_texts)
+ scores = np.empty(n_texts, dtype=np.float32)
+ start = 0
+ end = 0
+
+ while end < n_texts:
+ end += half_batch_size
+ batch_idx = slice(start, end)
+
+<<<<<<< HEAD
+=======
+ # Tokenize sentences
+>>>>>>> a24e0f2 (removed dependencies chek)
+ encoded_input = self.tokenizer(
+ source_texts[batch_idx] + ref_texts[batch_idx],
+ padding=True,
+ truncation=True,
+ return_tensors="pt",
+ )
+ encoded_input = {
+ key: value.to(self.config.device) for key, value in encoded_input.items()
+ }
+
+<<<<<<< HEAD
+ with torch.no_grad():
+ model_output = self.model(**encoded_input)
+
+ sent_embs = self.mean_pooling(model_output, encoded_input["attention_mask"])
+
+=======
+ # Calculate embeddings
+ with torch.no_grad():
+ model_output = self.model(**encoded_input)
+
+ # Perform pooling
+ sent_embs = self.mean_pooling(model_output, encoded_input["attention_mask"])
+
+ # Normalize embeddings
+>>>>>>> a24e0f2 (removed dependencies chek)
+ sent_embs = F.normalize(sent_embs, p=2, dim=1)
+
+ n_source_embs = len(sent_embs) // 2
+ scores[batch_idx] = (
+ (sent_embs[:n_source_embs] * sent_embs[n_source_embs:])
+ .sum(-1)
+ .cpu()
+ .detach()
+ .numpy()
+ )
+ start = end
+
+ return scores
+
+ def calculate(self, predictions: List[str], references: Optional[List[Union[str, List[str]]]] = None, original_texts: Optional[List[str]] = None) -> Dict[str, float]:
+ """
+ Calculate SentenceBERT semantic similarity.
+
+ Args:
+ predictions: List of predicted texts
+ references: List of reference texts
+ original_texts: List of original/source texts
+
+ Returns:
+ Dictionary with semantic similarity scores
+ """
+ self._initialize_model()
+
+ scores = {}
+
+<<<<<<< HEAD
+ if references is not None:
+ if isinstance(references[0], list):
+=======
+ # Similarity between predictions and references
+ if references is not None:
+ if isinstance(references[0], list):
+ # Handle multiple references - compute average similarity
+>>>>>>> a24e0f2 (removed dependencies chek)
+ ref_scores = []
+ for pred, ref_list in zip(predictions, references):
+ pred_list = [pred] * len(ref_list)
+ sim_scores = self._compute_similarity(pred_list, ref_list)
+ ref_scores.append(np.mean(sim_scores))
+ scores["sentbert_pred_ref"] = np.array(ref_scores)
+ else:
+ scores["sentbert_pred_ref"] = self._compute_similarity(predictions, references)
+
+<<<<<<< HEAD
+ if original_texts is not None:
+ scores["sentbert_pred_src"] = self._compute_similarity(predictions, original_texts)
+
+=======
+ # Similarity between predictions and original texts
+ if original_texts is not None:
+ scores["sentbert_pred_src"] = self._compute_similarity(predictions, original_texts)
+
+ # Aggregate if requested
+>>>>>>> a24e0f2 (removed dependencies chek)
+ if self.config.aggregate:
+ scores = {key: float(np.mean(value)) for key, value in scores.items()}
+
+ return scores
\ No newline at end of file
diff --git a/src/atgen/run_scripts/run_active_learning.py b/src/atgen/run_scripts/run_active_learning.py
index b22824a..a3747a1 100644
--- a/src/atgen/run_scripts/run_active_learning.py
+++ b/src/atgen/run_scripts/run_active_learning.py
@@ -9,7 +9,13 @@
from typing import Union
import logging
from atgen.utils.main_decorator import main_decorator
-from atgen.utils.constants import DEFAULT_CONFIG_NAME
+from atgen.utils.constants import (
+ DEFAULT_CONFIG_NAME,
+ UNLABELED_DATA_SPLIT_DEFAULT_NAME,
+ TEST_DATA_SPLIT_DEFAULT_NAME,
+ OUTPUT_FIELD_PURPOSE_TRAIN,
+ OUTPUT_FIELD_PURPOSE_TEST,
+)
log = logging.getLogger()
@@ -17,12 +23,16 @@
@main_decorator
def run_active_learning(config, workdir: Union[str, Path]):
from transformers import set_seed
- from datasets import concatenate_datasets
+ from datasets import concatenate_datasets, Dataset
from atgen.metrics.compute_metrics import compute_metrics
- from atgen.utils.data.load_data import load_data
- from atgen.utils.data.prepare_conversational_data import prepare_conversational_data
- from atgen.utils.data.maybe_get_few_shot_examples import maybe_get_few_shot_examples
+ from atgen.utils.data import (
+ load_data,
+ prepare_conversational_data,
+ maybe_get_few_shot_examples,
+ get_output_column_name_for_phase,
+ )
+ from atgen.utils.installers import install_spacy, install_nltk
from atgen.utils.load_model_tokenizer import load_model_tokenizer
from atgen.utils.prepare_model_for_training import prepare_model_for_training
from atgen.utils.training_utils import get_trainer
@@ -35,16 +45,22 @@ def run_active_learning(config, workdir: Union[str, Path]):
from atgen.utils.get_initial_labeled_data import (
get_initial_labeled_data_with_few_shot,
)
+ from atgen.strategies.base_strategy import BaseStrategy
from atgen.labellers.base_labeller import BaseLabeler
from atgen.utils.check_performance_metrics import (
check_performance_against_requirements,
)
+ # TODO Figure out how to stop downloading it every time
+ install_spacy()
+ install_nltk()
+
seed = config.seed
cache_dir = config.cache_dir
input_column_name = config.data.input_column_name
- output_column_name = config.data.output_column_name
dev_split_size = config.training.dev_split_size
+ output_column_name_train = config.data.train_output_column_name
+ output_column_name_test = config.data.test_output_column_name
model_name = config.model.checkpoint
@@ -62,7 +78,7 @@ def run_active_learning(config, workdir: Union[str, Path]):
# Initialize variables for tracking available metrics
available_metrics = {}
- metrics_availability_checked = False
+ is_metrics_availability_checked = False
has_test = (
config.data.test_split_name is not None and config.data.test_split_name != ""
@@ -89,14 +105,14 @@ def run_active_learning(config, workdir: Union[str, Path]):
log.info("Loading data.")
unlabeled_data = load_data(
data_config=config.data,
- split=config.data.unlabeled_data_split_name,
+ split=UNLABELED_DATA_SPLIT_DEFAULT_NAME,
cache_dir=config.cache_dir,
seed=seed,
)
if has_test:
test_data = load_data(
data_config=config.data,
- split=config.data.test_split_name,
+ split=TEST_DATA_SPLIT_DEFAULT_NAME,
cache_dir=config.cache_dir,
seed=seed,
)
@@ -107,7 +123,7 @@ def run_active_learning(config, workdir: Union[str, Path]):
)
log.info("Loading AL strategy.")
- al_strategy = get_strategy(
+ al_strategy: BaseStrategy = get_strategy(
config.al.strategy,
subsample_size=config.al.subsample_size,
unlabeled_pool=unlabeled_data[input_column_name],
@@ -124,7 +140,7 @@ def run_active_learning(config, workdir: Union[str, Path]):
# TODO: unsure whether need to log here since may be confusing for a human labeller
labeller: BaseLabeler = get_labeller(
config.labeller,
- output_column_name,
+ output_column_name=output_column_name_train,
cache_dir=cache_dir,
budget=budget,
workdir=workdir, # if labeller is a human
@@ -150,10 +166,10 @@ def run_active_learning(config, workdir: Union[str, Path]):
lambda x: x["id"] not in set(labeled_ids)
)
else:
- query_ids = al_strategy(
+ query_ids: list[str] = al_strategy(
model=model,
tokenizer=tokenizer,
- unlabeled_pool=unlabeled_data.remove_columns(output_column_name),
+ unlabeled_pool=unlabeled_data.remove_columns(output_column_name_train),
labeled_pool=None,
num_to_label=al_query_size,
batch_size=config.inference.batch_size,
@@ -185,7 +201,7 @@ def run_active_learning(config, workdir: Union[str, Path]):
labeled_data = unlabeled_data.select(range(0, 0))
labeled_ids = []
- unlabeled_data = prepare_conversational_data(
+ unlabeled_data: Dataset = prepare_conversational_data(
dataset=unlabeled_data,
data_config=config.data,
split="test",
@@ -194,16 +210,17 @@ def run_active_learning(config, workdir: Union[str, Path]):
)
if has_test:
- test_data = prepare_conversational_data(
- dataset=test_data,
- data_config=config.data,
- split="test",
- few_shot_examples=few_shot_examples,
- model_name=model_name,
- )
+ if not config.data.use_test_benchmark:
+ test_data: Dataset = prepare_conversational_data(
+ dataset=test_data,
+ data_config=config.data,
+ split="test",
+ few_shot_examples=few_shot_examples,
+ model_name=model_name,
+ )
# Evaluate the initial model before any training
if init_query_size_is_positive and config.al.evaluate_zero_iteration:
- generations = generate(
+ generations: list[str] = generate(
config.inference,
data=test_data,
model=model,
@@ -215,20 +232,20 @@ def run_active_learning(config, workdir: Union[str, Path]):
if os.path.exists(save_dir):
rmtree(save_dir)
- metrics = compute_metrics(
+ metrics: dict[str, float] = compute_metrics(
generated_texts=generations,
- reference_texts=test_data[output_column_name],
+ reference_texts=test_data[output_column_name_test],
original_texts=test_data[input_column_name],
config=config.evaluation,
cache_dir=cache_dir,
)
# Check required performance metrics
- is_performance_reached, metrics_availability_checked, available_metrics = (
+ is_performance_reached, is_metrics_availability_checked, available_metrics = (
check_performance_against_requirements(
metrics=metrics,
required_performance_dict=required_performance_dict,
- metrics_availability_checked=metrics_availability_checked,
+ is_metrics_availability_checked=is_metrics_availability_checked,
available_metrics=available_metrics,
)
)
@@ -284,7 +301,7 @@ def run_active_learning(config, workdir: Union[str, Path]):
# Set seed for reproducibility
set_seed(seed)
- trainer = get_trainer(
+ trainer: SFTTrainer = get_trainer(
config=config,
model=model,
tokenizer=tokenizer,
@@ -314,7 +331,7 @@ def run_active_learning(config, workdir: Union[str, Path]):
if dev_split_size > 0:
test_data = eval_data
else:
- generations = generate(
+ generations: list[str] = generate(
config.inference,
data=test_data,
model=model,
@@ -326,20 +343,20 @@ def run_active_learning(config, workdir: Union[str, Path]):
if os.path.exists(save_dir):
rmtree(save_dir)
- metrics = compute_metrics(
+ metrics: dict[str, float] = compute_metrics(
generated_texts=generations,
- reference_texts=test_data[output_column_name],
+ reference_texts=test_data[output_column_name_test],
original_texts=test_data[input_column_name],
config=config.evaluation,
cache_dir=cache_dir,
)
# Check required performance metrics
- is_performance_reached, metrics_availability_checked, available_metrics = (
+ is_performance_reached, is_metrics_availability_checked, available_metrics = (
check_performance_against_requirements(
metrics=metrics,
required_performance_dict=required_performance_dict,
- metrics_availability_checked=metrics_availability_checked,
+ is_metrics_availability_checked=is_metrics_availability_checked,
available_metrics=available_metrics,
)
)
@@ -357,22 +374,22 @@ def run_active_learning(config, workdir: Union[str, Path]):
# Make AL query for the next round if we have not run out of iterations
if al_iter != num_al_iterations + 1:
log.info(f"Making AL query at iteration {al_iter}.")
- query_ids = al_strategy(
+ query_ids: list[str] = al_strategy(
model=model,
tokenizer=tokenizer,
- unlabeled_pool=unlabeled_data.remove_columns(output_column_name),
+ unlabeled_pool=unlabeled_data.remove_columns(output_column_name_train),
labeled_pool=labeled_data,
num_to_label=al_query_size,
batch_size=config.inference.batch_size,
max_new_tokens=config.inference.max_new_tokens,
)
- query = unlabeled_data.filter(lambda x: x["id"] in query_ids)
- unlabeled_data = unlabeled_data.filter(lambda x: x["id"] not in query_ids)
- labeled_query = labeller(query)
+ query: Dataset = unlabeled_data.filter(lambda x: x["id"] in query_ids)
+ unlabeled_data: Dataset = unlabeled_data.filter(lambda x: x["id"] not in query_ids)
+ labeled_query: Dataset = labeller(query)
if labeller.is_out_of_budget:
log.info(f"Labeler ran out of budget at iteration {al_iter}.")
- labeled_data = concatenate_datasets([labeled_data, labeled_query])
+ labeled_data: Dataset = concatenate_datasets([labeled_data, labeled_query])
labeled_ids += query_ids
log.info(f"Saving labeled data at iteration #{al_iter}.")
diff --git a/src/atgen/strategies/base_strategy.py b/src/atgen/strategies/base_strategy.py
index 232fc66..eb96174 100644
--- a/src/atgen/strategies/base_strategy.py
+++ b/src/atgen/strategies/base_strategy.py
@@ -8,7 +8,7 @@
)
-class Strategy(ABC):
+class BaseStrategy(ABC):
def __init__(self, subsample_size: int | float = -1):
self.subsample_size = subsample_size
diff --git a/src/atgen/strategies/bleuvar.py b/src/atgen/strategies/bleuvar.py
index cd6e845..e9bccb6 100644
--- a/src/atgen/strategies/bleuvar.py
+++ b/src/atgen/strategies/bleuvar.py
@@ -1,11 +1,11 @@
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
import numpy as np
from datasets import Dataset
from transformers import set_seed, PreTrainedModel
-class BLEUVarStrategy(Strategy):
+class BLEUVarStrategy(BaseStrategy):
def __init__(self):
super().__init__()
diff --git a/src/atgen/strategies/dual.py b/src/atgen/strategies/dual.py
index a6ed8b1..0382fcf 100644
--- a/src/atgen/strategies/dual.py
+++ b/src/atgen/strategies/dual.py
@@ -1,7 +1,7 @@
from datasets import Dataset
from transformers import PreTrainedModel
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
from .idds import idds_sampling
from .bleuvar import bleuvar
from .random_strategy import random_strategy
@@ -11,7 +11,7 @@
"""
-class DualStrategy(Strategy):
+class DualStrategy(BaseStrategy):
def __init__(self, seed=None):
super().__init__()
self.seed = seed
diff --git a/src/atgen/strategies/graph_cut.py b/src/atgen/strategies/graph_cut.py
index c6120ce..d03ebc9 100644
--- a/src/atgen/strategies/graph_cut.py
+++ b/src/atgen/strategies/graph_cut.py
@@ -5,7 +5,7 @@
from submodlib import GraphCutFunction
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
from ..utils.get_embeddings import get_embeddings
"""
@@ -15,7 +15,7 @@
"""
-class GraphCutStrategy(Strategy):
+class GraphCutStrategy(BaseStrategy):
def __init__(
self,
unlabeled_pool: list[str],
diff --git a/src/atgen/strategies/hadas.py b/src/atgen/strategies/hadas.py
index f98a215..a18ca2b 100644
--- a/src/atgen/strategies/hadas.py
+++ b/src/atgen/strategies/hadas.py
@@ -1,8 +1,8 @@
-from typing import Optional
+from typing import Optional, Union
import numpy as np
from evaluate import EvaluationModule, load
from scipy.spatial.distance import jensenshannon
-from omegaconf import DictConfig
+from omegaconf import DictConfig, ListConfig
import logging
from datasets import Dataset
@@ -18,16 +18,16 @@
)
from .unieval import SumEvaluator, convert_to_json
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
from ..utils.generate import generate
-from ..utils.data.prepare_conversational_data import prepare_conversational_data
+from ..utils.data import prepare_conversational_data, get_output_column_name_for_phase
from ..utils.constants import MESSAGES_COLUMN_NAME
log = logging.getLogger()
-class HadasStrategy(Strategy):
+class HadasStrategy(BaseStrategy):
def __init__(
self,
subsample_size: int | float = -1,
@@ -81,7 +81,7 @@ def __call__(
unlabeled_pool=unlabeled_pool,
labeled_pool=labeled_pool,
input_column_name=self.data_config.input_column_name,
- output_column_name=self.data_config.output_column_name,
+ output_column_name=self.data_config.train_output_column_name,
num_to_label=num_to_label,
inference_config=self.inference_config,
model_config=self.model_config,
@@ -221,13 +221,13 @@ def hadas(
h_halu = (U_unlabeled @ weights).flatten().numpy()
U_labeled = hallucination_distribution(
- entailment_model,
- entailment_tokenizer,
- unieval,
- bertscore,
- labeled_pool[input_column_name],
- labeled_pool[output_column_name],
- inference_config.batch_size,
+ entailment_model=entailment_model,
+ entailment_tokenizer=entailment_tokenizer,
+ unieval=unieval,
+ bertscore=bertscore,
+ documents=labeled_pool[input_column_name],
+ summaries=labeled_pool[output_column_name],
+ batch_size=inference_config.batch_size,
)
h_div = np.array(
[
diff --git a/src/atgen/strategies/huds.py b/src/atgen/strategies/huds.py
index ccf9449..473b729 100644
--- a/src/atgen/strategies/huds.py
+++ b/src/atgen/strategies/huds.py
@@ -13,13 +13,13 @@
PreTrainedTokenizerBase,
)
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
from ..utils.get_embeddings import get_embeddings
from ..utils.constants import MESSAGES_COLUMN_NAME
from ..utils.data.prepare_conversational_data import prepare_conversational_data
-class HudsStrategy(Strategy):
+class HudsStrategy(BaseStrategy):
def __init__(
self,
unlabeled_pool: list[str],
diff --git a/src/atgen/strategies/idds.py b/src/atgen/strategies/idds.py
index 626cea3..2dcf987 100644
--- a/src/atgen/strategies/idds.py
+++ b/src/atgen/strategies/idds.py
@@ -1,4 +1,4 @@
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
from datasets import Dataset
from transformers import PreTrainedModel
@@ -10,7 +10,7 @@
import logging
-class IDDSStrategy(Strategy):
+class IDDSStrategy(BaseStrategy):
def __init__(self, seed=None):
super().__init__()
self.seed = seed
diff --git a/src/atgen/strategies/nsp.py b/src/atgen/strategies/nsp.py
index 3c6438b..fdcaba5 100644
--- a/src/atgen/strategies/nsp.py
+++ b/src/atgen/strategies/nsp.py
@@ -1,4 +1,4 @@
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
from typing import Union
import numpy as np
@@ -6,7 +6,7 @@
from transformers import PreTrainedModel
-class NSPStrategy(Strategy):
+class NSPStrategy(BaseStrategy):
def __init__(self):
super().__init__()
diff --git a/src/atgen/strategies/random_strategy.py b/src/atgen/strategies/random_strategy.py
index 019a0f9..556391f 100644
--- a/src/atgen/strategies/random_strategy.py
+++ b/src/atgen/strategies/random_strategy.py
@@ -1,7 +1,7 @@
import numpy as np
from datasets import Dataset
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
def random_strategy(
@@ -15,7 +15,7 @@ def random_strategy(
return ids[:num_to_label]
-class RandomStrategy(Strategy):
+class RandomStrategy(BaseStrategy):
def __init__(self, subsample_size: int = -1, seed: int = 42):
super().__init__()
self.seed = seed
diff --git a/src/atgen/strategies/te_delfy.py b/src/atgen/strategies/te_delfy.py
index 2e98f77..08c1da0 100644
--- a/src/atgen/strategies/te_delfy.py
+++ b/src/atgen/strategies/te_delfy.py
@@ -17,13 +17,13 @@
PreTrainedTokenizerBase,
)
-from .base_strategy import Strategy
+from .base_strategy import BaseStrategy
log = logging.getLogger()
-class TeDelfyStrategy(Strategy):
+class TeDelfyStrategy(BaseStrategy):
def __init__(self, subsample_size: int = -1, inference_config: DictConfig = None):
super().__init__(subsample_size)
self.random_init = True
diff --git a/src/atgen/utils/check_performance_metrics.py b/src/atgen/utils/check_performance_metrics.py
index 8a8b0a8..9702cc0 100644
--- a/src/atgen/utils/check_performance_metrics.py
+++ b/src/atgen/utils/check_performance_metrics.py
@@ -4,7 +4,7 @@
def check_performance_against_requirements(
- metrics, required_performance_dict, metrics_availability_checked, available_metrics
+ metrics, required_performance_dict, is_metrics_availability_checked, available_metrics
):
"""
Check if the computed metrics meet the required performance thresholds.
@@ -12,17 +12,17 @@ def check_performance_against_requirements(
Args:
metrics (dict): The computed metrics from the current evaluation
required_performance_dict (dict): Dictionary of required metric thresholds
- metrics_availability_checked (bool): Whether metrics availability has been checked
+ is_metrics_availability_checked (bool): Whether metrics availability has been checked
available_metrics (dict): Previously identified available metrics and their thresholds
Returns:
- tuple: (is_performance_reached, metrics_availability_checked, available_metrics)
+ tuple: (is_performance_reached, is_metrics_availability_checked, available_metrics)
"""
is_performance_reached = False
if required_performance_dict is not None:
# Only check which metrics are available on the first iteration with valid metrics
- if not metrics_availability_checked:
+ if not is_metrics_availability_checked:
# Determine which required metrics are available in computed metrics
available_metrics = {
metric: threshold
@@ -44,7 +44,7 @@ def check_performance_against_requirements(
"None of the required metrics are available. Cannot evaluate required performance."
)
- metrics_availability_checked = True
+ is_metrics_availability_checked = True
# Check if required performance is reached based on available metrics
if available_metrics:
@@ -75,4 +75,4 @@ def check_performance_against_requirements(
else:
is_performance_reached = False
- return is_performance_reached, metrics_availability_checked, available_metrics
+ return is_performance_reached, is_metrics_availability_checked, available_metrics
diff --git a/src/atgen/utils/constants.py b/src/atgen/utils/constants.py
index 2bfff31..95fd7e9 100644
--- a/src/atgen/utils/constants.py
+++ b/src/atgen/utils/constants.py
@@ -6,7 +6,10 @@
DEFAULT_TEMPERATURE = 0.6
DEFAULT_TOP_P = 1
-DEEPSEEK_R1_END_REASONING_TOKEN = ""
+REASONING_END_TOKEN = ""
APPROX_NUM_SYSTEM_TOKENS_PER_MESSAGE = 10
DEFAULT_CONFIG_NAME = "base"
+
+OUTPUT_FIELD_PURPOSE_TRAIN = "train"
+OUTPUT_FIELD_PURPOSE_TEST = "test"
diff --git a/src/atgen/utils/data/__init__.py b/src/atgen/utils/data/__init__.py
index 454678c..33c7776 100644
--- a/src/atgen/utils/data/__init__.py
+++ b/src/atgen/utils/data/__init__.py
@@ -1,3 +1,5 @@
+from .get_output_column_name_for_phase import get_output_column_name_for_phase
+from .get_preprocess_function import get_preprocess_function
from .prepare_conversational_data import prepare_conversational_data
from .load_data import load_data
from .maybe_get_few_shot_examples import maybe_get_few_shot_examples
diff --git a/src/atgen/utils/data/get_output_column_name.py b/src/atgen/utils/data/get_output_column_name.py
new file mode 100644
index 0000000..c3947f1
--- /dev/null
+++ b/src/atgen/utils/data/get_output_column_name.py
@@ -0,0 +1,19 @@
+from typing import Union
+from omegaconf import DictConfig, ListConfig
+
+from ..constants import OUTPUT_FIELD_PURPOSE_TRAIN
+
+
+def get_output_column_name(
+ output_column_name: Union[DictConfig, ListConfig, dict[str, Union[str, list[str]]], list[str], str],
+ purpose: str = OUTPUT_FIELD_PURPOSE_TRAIN,
+) -> str:
+ if isinstance(output_column_name, (list, ListConfig)):
+ return '_'.join(output_column_name)
+ elif isinstance(output_column_name, (dict, DictConfig)):
+ return '_'.join(output_column_name[purpose])
+ elif isinstance(output_column_name, str):
+ return output_column_name
+ else:
+ raise NotImplementedError(f"Unexpected type {type(output_column_name)} of the output column name.")
+
\ No newline at end of file
diff --git a/src/atgen/utils/data/get_output_column_name_for_phase.py b/src/atgen/utils/data/get_output_column_name_for_phase.py
new file mode 100644
index 0000000..84a34e2
--- /dev/null
+++ b/src/atgen/utils/data/get_output_column_name_for_phase.py
@@ -0,0 +1,22 @@
+from typing import Union
+from omegaconf import DictConfig, ListConfig
+
+from ..constants import OUTPUT_FIELD_PURPOSE_TRAIN
+
+
+def get_output_column_name_for_phase(
+ output_column_name: Union[
+ DictConfig, ListConfig, dict[str, Union[str, list[str]]], list[str], str
+ ],
+ purpose: str = OUTPUT_FIELD_PURPOSE_TRAIN,
+) -> str:
+ if isinstance(output_column_name, (list, ListConfig)):
+ return "_".join(output_column_name)
+ elif isinstance(output_column_name, (dict, DictConfig)):
+ return "_".join(output_column_name[purpose])
+ elif isinstance(output_column_name, str):
+ return output_column_name
+ else:
+ raise NotImplementedError(
+ f"Unexpected type {type(output_column_name)} of the output column name."
+ )
diff --git a/src/atgen/utils/data/get_preprocess_function.py b/src/atgen/utils/data/get_preprocess_function.py
index ba31b8b..0110654 100644
--- a/src/atgen/utils/data/get_preprocess_function.py
+++ b/src/atgen/utils/data/get_preprocess_function.py
@@ -1,6 +1,6 @@
-from omegaconf import DictConfig
+from typing import Optional
-from ..constants import MESSAGES_COLUMN_NAME
+from ..constants import MESSAGES_COLUMN_NAME, OUTPUT_FIELD_PURPOSE_TRAIN
def get_preprocess_function(
@@ -11,6 +11,7 @@ def get_preprocess_function(
is_in_conversational_format: bool = False,
input_column_name: str = "input",
output_column_name: str = "output",
+ assistant_response_start: Optional[str] = None,
):
"""
Creates and returns an appropriate preprocessing function based on:
@@ -36,9 +37,7 @@ def preprocess_fn(instance):
# Copy the first few-shot message and prepend system prompt
first_fs = few_shot_messages[0].copy()
first_fs["content"] = system_prompt + "\n\n" + first_fs["content"]
- messages.append(first_fs)
- # Add remaining few-shot messages
- messages.extend(few_shot_messages[1:])
+ messages = [first_fs] + few_shot_messages[1:]
else:
# No few-shot messages
if is_in_conversational_format:
@@ -47,6 +46,9 @@ def preprocess_fn(instance):
conv_messages[0]["content"] = (
system_prompt + "\n\n" + conv_messages[0]["content"]
)
+ conv_messages = _maybe_add_assistant_response_start(
+ conv_messages, assistant_response_start
+ )
return {MESSAGES_COLUMN_NAME: conv_messages}
else:
# Add user message with system prompt
@@ -66,10 +68,11 @@ def preprocess_fn(instance):
if is_in_conversational_format:
# Return the messages directly if in conversational format
- return {
- MESSAGES_COLUMN_NAME: few_shot_messages
- + instance[input_column_name]
- }
+ messages = few_shot_messages + instance[input_column_name]
+ messages = _maybe_add_assistant_response_start(
+ messages, assistant_response_start
+ )
+ return {MESSAGES_COLUMN_NAME: messages}
else:
# Add user message
messages.append(
@@ -79,11 +82,23 @@ def preprocess_fn(instance):
# Handle different split types
if is_in_conversational_format:
# For conversational format, we've already handled this above
- return {MESSAGES_COLUMN_NAME: messages + instance[input_column_name]}
+ messages += instance[input_column_name]
+ messages = _maybe_add_assistant_response_start(
+ messages, assistant_response_start
+ )
+ return {MESSAGES_COLUMN_NAME: messages}
elif split == "train":
# For training, add the assistant's response
+ if assistant_response_start:
+ assistant_message = (
+ assistant_response_start + instance[output_column_name]
+ )
+ else:
+ assistant_message = instance[output_column_name]
+ messages.append({"role": "assistant", "content": assistant_message})
+ elif split == "test" and assistant_response_start:
messages.append(
- {"role": "assistant", "content": instance[output_column_name]}
+ {"role": "assistant", "content": assistant_response_start}
)
return {MESSAGES_COLUMN_NAME: messages}
@@ -104,7 +119,11 @@ def preprocess_fn(instance):
# Handle different formats
if is_in_conversational_format:
# For conversational format, append input messages to existing ones
- return {MESSAGES_COLUMN_NAME: messages + instance[input_column_name]}
+ messages += instance[input_column_name]
+ messages = _maybe_add_assistant_response_start(
+ messages, assistant_response_start
+ )
+ return {MESSAGES_COLUMN_NAME: messages}
else:
# Add user message
messages.append(
@@ -113,10 +132,26 @@ def preprocess_fn(instance):
# For training, add the assistant's response
if split == "train":
+ if assistant_response_start:
+ assistant_message = (
+ assistant_response_start + instance[output_column_name]
+ )
+ else:
+ assistant_message = instance[output_column_name]
+ messages.append({"role": "assistant", "content": assistant_message})
+ elif split == "test" and assistant_response_start:
messages.append(
- {"role": "assistant", "content": instance[output_column_name]}
+ {"role": "assistant", "content": assistant_response_start}
)
return {MESSAGES_COLUMN_NAME: messages}
return preprocess_fn
+
+
+def _maybe_add_assistant_response_start(
+ messages: list[dict[str, str]], assistant_response_start: Optional[str]
+) -> list[dict[str, str]]:
+ if assistant_response_start:
+ messages[-1]["content"] = assistant_response_start + messages[-1]["content"]
+ return messages
diff --git a/src/atgen/utils/data/load_data.py b/src/atgen/utils/data/load_data.py
index 0351b4b..dbc5a5b 100644
--- a/src/atgen/utils/data/load_data.py
+++ b/src/atgen/utils/data/load_data.py
@@ -2,7 +2,10 @@
from typing import Union
from datasets import load_dataset, load_from_disk, Dataset, DatasetDict
-from omegaconf import DictConfig
+from omegaconf import DictConfig, ListConfig
+
+from .get_output_column_name_for_phase import get_output_column_name_for_phase
+from ..constants import OUTPUT_FIELD_PURPOSE_TRAIN, OUTPUT_FIELD_PURPOSE_TEST
def _fetch_dataset(
@@ -10,8 +13,11 @@ def _fetch_dataset(
subset_name: str,
fetch_kwargs: dict | DictConfig,
) -> Dataset:
+ # Load a subset of a dataset from HuggingFace
+ if isinstance(dataset_name_or_path, (list, ListConfig)):
+ dataset = load_dataset(*dataset_name_or_path, **fetch_kwargs)
# Load local dataset
- if os.path.exists(dataset_name_or_path):
+ elif os.path.exists(dataset_name_or_path):
# Load a saved on disk dataset
if os.path.isdir(dataset_name_or_path):
# Remove `cache_dir` from fetch_kwargs
@@ -28,8 +34,6 @@ def _fetch_dataset(
f"Unexpected format {dataset_name_or_path.split('.')[-1]} of the dataset. Supported formats: csv, json."
)
# Load dataset from HuggingFace
- elif isinstance(dataset_name_or_path, list):
- dataset = load_dataset(*dataset_name_or_path, **fetch_kwargs)
else:
dataset = load_dataset(dataset_name_or_path, **fetch_kwargs)
@@ -51,10 +55,50 @@ def _take_subset(dataset_subset: Dataset, size: int, seed: int) -> Dataset:
return dataset_subset
dataset_subset = dataset_subset.shuffle(seed=seed)
dataset_subset = dataset_subset.select(range(size))
- dataset_subset = dataset_subset.remove_columns(["id"]).add_column("id", list(range(len(dataset_subset))))
+ dataset_subset = dataset_subset.remove_columns(["id"]).add_column(
+ "id", list(range(len(dataset_subset)))
+ )
return dataset_subset
+def _preprocess_multicolumn_labels_if_needed(
+ dataset: Dataset,
+ output_column_names: Union[
+ DictConfig, ListConfig, dict[str, Union[str, list[str]]], list[str], str
+ ],
+ data_config: DictConfig,
+ phase: str = OUTPUT_FIELD_PURPOSE_TRAIN,
+) -> Dataset:
+ if isinstance(output_column_names, (list, ListConfig)):
+ # Get preprocessed column name
+ if phase == OUTPUT_FIELD_PURPOSE_TRAIN:
+ new_column_name = data_config.train_output_column_name
+ elif phase == OUTPUT_FIELD_PURPOSE_TEST:
+ new_column_name = data_config.test_output_column_name
+ else:
+ raise NotImplementedError(f"Unexpected phase {phase}")
+ # Get preprocessed column values
+ values = []
+ for inst in dataset:
+ for col_name in output_column_names:
+ inst = inst[col_name]
+ values.append(inst)
+ dataset = dataset.add_column(new_column_name, values)
+ elif isinstance(output_column_names, (dict, DictConfig)):
+ for phase, column_name in output_column_names.items():
+ dataset = _preprocess_multicolumn_labels_if_needed(
+ dataset, column_name, data_config, phase
+ )
+ # Nothing to preprocess in this case
+ elif isinstance(output_column_names, str):
+ pass
+ else:
+ raise NotImplementedError(
+ f"Unexpected type {type(output_column_names)} of the output column names."
+ )
+ return dataset
+
+
def load_data(
data_config: DictConfig,
split: str,
@@ -62,21 +106,26 @@ def load_data(
seed: int,
) -> Dataset:
if split == "train":
- subset_name = data_config.get("train_subset_name", split)
+ subset_name = data_config.get("train_split_name", split)
subset_size = data_config.get("train_subset_size")
elif split == "test":
- subset_name = data_config.get("test_subset_name", split)
+ subset_name = data_config.get("test_split_name", split)
subset_size = data_config.get("test_subset_size")
else:
raise NotImplementedError(
f"Unexpected split {split}; Please specify either `train` or `test`."
)
dataset = _fetch_dataset(
- data_config.dataset,
- subset_name,
- dict(data_config.fetch_kwargs, cache_dir=cache_dir),
+ dataset_name_or_path=data_config.dataset,
+ subset_name=subset_name,
+ fetch_kwargs=dict(data_config.fetch_kwargs, cache_dir=cache_dir),
+ )
+ dataset = _preprocess_multicolumn_labels_if_needed(
+ dataset=dataset,
+ output_column_names=data_config.output_column_name,
+ data_config=data_config,
+ phase=split,
)
-
# Add `id` column to the dataset (practical use) or to train subset (benchmarking)
dataset = _add_id_column(dataset)
if subset_size is not None:
diff --git a/src/atgen/utils/data/prepare_conversational_data.py b/src/atgen/utils/data/prepare_conversational_data.py
index 86f1f10..432d9c3 100644
--- a/src/atgen/utils/data/prepare_conversational_data.py
+++ b/src/atgen/utils/data/prepare_conversational_data.py
@@ -1,9 +1,9 @@
from itertools import chain
from datasets import Dataset
from omegaconf import DictConfig
-from transformers import PreTrainedTokenizer
from .get_preprocess_function import get_preprocess_function
+from .get_output_column_name_for_phase import get_output_column_name_for_phase
def prepare_conversational_data(
@@ -14,25 +14,16 @@ def prepare_conversational_data(
model_name: str = "kek",
) -> Dataset:
input_column_name = data_config.input_column_name
- output_column_name = data_config.output_column_name
+ output_column_name = get_output_column_name_for_phase(
+ output_column_name=data_config.output_column_name, purpose=split
+ )
- if not few_shot_examples:
- few_shot_messages = []
- else:
- few_shot_messages = list(
- chain.from_iterable(
- [
- [
- {"role": "user", "content": fs_input},
- {"role": "assistant", "content": fs_output},
- ]
- for (fs_input, fs_output) in zip(
- few_shot_examples[input_column_name],
- few_shot_examples[output_column_name],
- )
- ]
- )
- )
+ few_shot_messages = []
+ if few_shot_examples:
+ for (fs_input, fs_output) in zip(few_shot_examples[input_column_name], few_shot_examples[output_column_name]):
+ few_shot_messages.append({"role": "user", "content": fs_input})
+ response_start = data_config.assistant_response_start if data_config.assistant_response_start else ""
+ few_shot_messages.append({"role": "assistant", "content": response_start + fs_output})
# Get appropriate preprocessing function based on all parameters
preprocess_fn = get_preprocess_function(
model_name=model_name,
@@ -42,8 +33,8 @@ def prepare_conversational_data(
is_in_conversational_format=data_config.is_in_conversational_format,
input_column_name=input_column_name,
output_column_name=output_column_name,
+ assistant_response_start=data_config.assistant_response_start,
)
-
dataset = dataset.map(
preprocess_fn,
batched=False,
diff --git a/src/atgen/utils/generate.py b/src/atgen/utils/generate.py
index 95a86c0..ef080c2 100644
--- a/src/atgen/utils/generate.py
+++ b/src/atgen/utils/generate.py
@@ -8,7 +8,7 @@
from omegaconf import DictConfig
from torch.utils.data import DataLoader
-from torch import bfloat16, float16, cuda, no_grad, LongTensor
+from torch import bfloat16, cuda
from tqdm import tqdm
from transformers import PreTrainedModel, PreTrainedTokenizer
from vllm import LLM
@@ -21,15 +21,19 @@
DEFAULT_TEMPERATURE,
DEFAULT_TOP_P,
MESSAGES_COLUMN_NAME,
- DEEPSEEK_R1_END_REASONING_TOKEN,
)
from .training_utils import _get_data_collator
from .find_response_token_ids_in_text import find_response_token_ids_in_text
+from .post_process_generations import post_process_generations
log = logging.getLogger()
+VLLM_FRAMEWORK = "vllm"
+SGLANG_FRAMEWORK = "sglang"
+TRANSFORMERS_FRAMEWORK = "transformers"
+
def generate_vllm(
inference_config: DictConfig,
@@ -72,12 +76,22 @@ def generate_vllm(
gc.collect()
cuda.empty_cache()
- params = SamplingParams(
+ sampling_params = SamplingParams(
temperature=inference_config.get("temperature", DEFAULT_TEMPERATURE),
seed=42, # TODO: make arbitrary
max_tokens=inference_config.max_new_tokens,
top_p=inference_config.get("top_p", DEFAULT_TOP_P),
)
+ if data_config.assistant_response_start:
+ generation_params = {
+ "add_generation_prompt": False,
+ "continue_final_message": True,
+ }
+ else:
+ generation_params = {
+ "add_generation_prompt": True,
+ "continue_final_message": False,
+ }
generations = []
num_batches = ceil(len(data) / inference_config.batch_size)
@@ -85,17 +99,18 @@ def generate_vllm(
batch = data[
i * inference_config.batch_size : (i + 1) * inference_config.batch_size
][MESSAGES_COLUMN_NAME]
- out = llm_runner.chat(batch, params, use_tqdm=False)
+ out = llm_runner.chat(
+ batch, sampling_params, use_tqdm=False, **generation_params
+ )
batch_generations = [x.outputs[0].text for x in out]
generations += batch_generations
- # Remove reasoning tokens from DeepSeek-R1
- if "deepseek-r1" in llm_runner.llm_engine.model_config.model:
- for i, generation in enumerate(generations):
- generations[i] = DEEPSEEK_R1_END_REASONING_TOKEN.join(
- generation.split(DEEPSEEK_R1_END_REASONING_TOKEN)[1:]
- ).strip()
-
+ generations = post_process_generations(
+ generations=generations,
+ data_config=data_config,
+ model_name=llm_runner.llm_engine.model_config.model,
+ framework=VLLM_FRAMEWORK
+ )
if delete_vllm_after_inference:
del llm_runner
gc.collect()
@@ -103,14 +118,6 @@ def generate_vllm(
return generations
-def generate_sglang() -> list[str]:
- """
- Function for generating with the SGLang framework.
- Requires either model + tokenizer or the path to the saved model and tokenizer.
- """
- pass
-
-
def generate_sglang(
inference_config: DictConfig,
data: Dataset,
@@ -198,6 +205,12 @@ def generate_response(s, messages):
batch_generations = [result.outputs["response"] for result in batch_results]
generations += batch_generations
+ generations = post_process_generations(
+ generations=generations,
+ data_config=data_config,
+ model_name=model_path,
+ framework=SGLANG_FRAMEWORK
+ )
# Clean up
engine.shutdown()
return generations
@@ -216,7 +229,7 @@ def generate_transformers(
data = data.map(
tokenize_conversational_example,
batched=False,
- fn_kwargs={"tokenizer": tokenizer},
+ fn_kwargs={"tokenizer": tokenizer, "data_config": data_config},
)
data_collator = _get_data_collator(
@@ -237,31 +250,36 @@ def generate_transformers(
FastLanguageModel.for_inference(model)
generations = []
- with no_grad():
- for batch in tqdm(dataloader):
- try:
- out = model.generate(
- batch["input_ids"].to(model.device),
- attention_mask=batch["attention_mask"].to(model.device),
- max_new_tokens=inference_config.max_new_tokens,
- temperature=inference_config.get(
- "temperature", DEFAULT_TEMPERATURE
- ),
- top_p=inference_config.get("top_p", DEFAULT_TOP_P),
- return_dict_in_generate=True,
- output_scores=True,
+ for batch in tqdm(dataloader):
+ try:
+ out = model.generate(
+ batch["input_ids"].to(model.device),
+ attention_mask=batch["attention_mask"].to(model.device),
+ max_new_tokens=inference_config.max_new_tokens,
+ temperature=inference_config.get(
+ "temperature", DEFAULT_TEMPERATURE
+ ),
+ top_p=inference_config.get("top_p", DEFAULT_TOP_P),
+ return_dict_in_generate=True,
+ output_scores=True,
+ )
+ outputs_only = []
+ for output in out.sequences:
+ seq_only = find_response_token_ids_in_text(
+ output.tolist(), data_collator.response_token_ids
)
- outputs_only = []
- for output in out.sequences:
- seq_only = find_response_token_ids_in_text(
- output.tolist(), data_collator.response_token_ids
- )
- outputs_only.append(tokenizer.decode(seq_only, True).strip())
- generations += outputs_only
- except Exception as e:
- print(f"Error in model.generate: {e}")
- # Add empty strings for this batch to maintain alignment with input data
- generations += [""] * len(batch["input_ids"])
+ outputs_only.append(tokenizer.decode(seq_only, True).strip())
+ generations += outputs_only
+ except Exception as e:
+ print(f"Error in model.generate: {e}")
+ # Add empty strings for this batch to maintain alignment with input data
+ generations += [""] * len(batch["input_ids"])
+ generations = post_process_generations(
+ generations=generations,
+ data_config=data_config,
+ model_name=model.name_or_path,
+ framework=TRANSFORMERS_FRAMEWORK
+ )
return generations
@@ -276,7 +294,7 @@ def generate(
**kwargs,
) -> list[str]:
framework = inference_config.framework
- if framework == "vllm":
+ if framework == VLLM_FRAMEWORK:
return generate_vllm(
inference_config=inference_config,
data=data,
@@ -286,7 +304,7 @@ def generate(
data_config=data_config,
**kwargs,
)
- elif framework == "transformers":
+ elif framework == TRANSFORMERS_FRAMEWORK:
return generate_transformers(
inference_config=inference_config,
data=data,
@@ -297,7 +315,7 @@ def generate(
model_config=model_config,
**kwargs,
)
- elif framework == "sglang":
+ elif framework == SGLANG_FRAMEWORK:
return generate_sglang(
inference_config=inference_config,
data=data,
@@ -313,8 +331,11 @@ def generate(
def tokenize_conversational_example(
- example: dict[str, Any], tokenizer: PreTrainedTokenizer
+ example: dict[str, Any], tokenizer: PreTrainedTokenizer, data_config: DictConfig
) -> dict[str, list[int]]:
- input_ids = tokenizer.apply_chat_template(example["messages"])
+ if data_config.assistant_response_start:
+ input_ids = tokenizer.apply_chat_template(example["messages"], continue_final_message=True)
+ else:
+ input_ids = tokenizer.apply_chat_template(example["messages"], add_generation_prompt=True)
attention_mask = [1 for _ in range(len(input_ids))]
return {"input_ids": input_ids, "attention_mask": attention_mask}
diff --git a/src/atgen/utils/installers.py b/src/atgen/utils/installers.py
new file mode 100644
index 0000000..b540fb9
--- /dev/null
+++ b/src/atgen/utils/installers.py
@@ -0,0 +1,10 @@
+import subprocess
+
+
+def install_spacy():
+ subprocess.run("python -m spacy download en_core_web_sm", shell=True)
+
+def install_nltk():
+ import nltk
+ nltk.download('punkt')
+ nltk.download('punkt_tab')
diff --git a/src/atgen/utils/load_model_tokenizer.py b/src/atgen/utils/load_model_tokenizer.py
index 2711240..f33a5ac 100644
--- a/src/atgen/utils/load_model_tokenizer.py
+++ b/src/atgen/utils/load_model_tokenizer.py
@@ -53,6 +53,7 @@ def load_model_tokenizer(
trust_remote_code=True,
)
tokenizer.model_max_length = model_config.model_max_length
+ tokenizer.padding_side = "left"
else:
kwargs = {
"cache_dir": cache_dir,
@@ -65,8 +66,7 @@ def load_model_tokenizer(
model_config.checkpoint,
model_max_length=model_config.model_max_length,
cache_dir=cache_dir,
- # TODO: choose automatically
- # padding_side="left",
+ padding_side="left"
)
if tokenizer.pad_token is None:
diff --git a/src/atgen/utils/post_process_generations.py b/src/atgen/utils/post_process_generations.py
new file mode 100644
index 0000000..fd5425a
--- /dev/null
+++ b/src/atgen/utils/post_process_generations.py
@@ -0,0 +1,21 @@
+from typing import Optional, Literal
+from omegaconf import DictConfig
+
+from .constants import REASONING_END_TOKEN
+
+def post_process_generations(
+ generations: list[str],
+ data_config: DictConfig,
+ model_name: Optional[str] = None,
+ framework: Literal["vllm", "sglang", "transformers"] = "vllm"
+) -> list[str]:
+ # Remove assistant response start from transformers generations
+ if framework == "transformers" and (ass_resp_start := data_config.assistant_response_start):
+ return [ass_resp_start.join(gen.split(ass_resp_start)[1:]).strip() for gen in generations]
+ elif model_name and "deepseek-r1" in model_name:
+ return [_remove_thinking_part(gen) for gen in generations]
+ return generations
+
+
+def _remove_thinking_part(text: str) -> str:
+ return REASONING_END_TOKEN.join(text.split(REASONING_END_TOKEN)[1:]).strip()
diff --git a/src/atgen/utils/resolvers.py b/src/atgen/utils/resolvers.py
index 629b453..983cfce 100644
--- a/src/atgen/utils/resolvers.py
+++ b/src/atgen/utils/resolvers.py
@@ -2,10 +2,9 @@
Custom Hydra resolvers for configuration calculations
"""
-from typing import Any
+from typing import Union
-from hydra.core.config_store import ConfigStore
-from omegaconf import OmegaConf
+from omegaconf import OmegaConf, DictConfig, ListConfig
def multiply_with_few_shot(input_max_length: int, few_shot_count: int) -> int:
@@ -28,15 +27,15 @@ def to_string(model_name: str):
"""
return model_name.replace("/", "__")
-
def register_resolvers() -> None:
"""Register all custom resolvers with OmegaConf"""
# Register resolvers only if they are not already registered
resolvers_to_register = {
"multiply_with_few_shot": multiply_with_few_shot,
"to_string": to_string,
+ "get_output_column_name_for_phase": get_output_column_name_for_phase,
}
-
+
for name, resolver_fn in resolvers_to_register.items():
try:
OmegaConf.register_new_resolver(name, resolver_fn)
diff --git a/src/atgen/utils/training_utils.py b/src/atgen/utils/training_utils.py
index 88c2d55..1646d3f 100644
--- a/src/atgen/utils/training_utils.py
+++ b/src/atgen/utils/training_utils.py
@@ -16,9 +16,9 @@
)
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM
from trl.data_utils import (
- is_conversational,
- maybe_apply_chat_template,
- maybe_convert_to_chatml
+ is_conversational,
+ maybe_apply_chat_template,
+ maybe_convert_to_chatml,
)
from trl.trainer.utils import ConstantLengthDataset
from accelerate import PartialState
@@ -87,7 +87,7 @@ class DataCollatorForLastCompletionOnlyLM(DataCollatorForCompletionOnlyLM):
Data collator that extends DataCollatorForCompletionOnlyLM to only train on the last assistant message.
It ensures that only the final assistant response will contribute to the loss, while all previous messages
(including earlier assistant responses) are masked with the ignore_index.
-
+
"""
def torch_call(self, examples):
@@ -241,8 +241,12 @@ def _get_train_eval_datasets(
batched=True,
fn_kwargs={"tokenizer": tokenizer, "data_collator": data_collator},
).filter(lambda x: x[TEXT_FIELD] != "")
- logger.warning(f"Truncated {orig_train_data_len - len(train_data)} examples from train set.")
- logger.warning(f"Truncated {orig_eval_data_len - len(eval_data)} examples from eval set.")
+ logger.warning(
+ f"Truncated {orig_train_data_len - len(train_data)} examples from train set."
+ )
+ logger.warning(
+ f"Truncated {orig_eval_data_len - len(eval_data)} examples from eval set."
+ )
return train_data, eval_data
@@ -269,6 +273,7 @@ def _formatting_fn(examples, tokenizer: PreTrainedTokenizerFast):
add_generation_prompt=False,
)
+
def get_trainer(
config: DictConfig,
model: PreTrainedModel | PeftModel,
@@ -298,7 +303,7 @@ def get_trainer(
tokenizer=tokenizer,
data_collator=data_collator,
)
-
+
return SFTTrainer(
model=model,
processing_class=tokenizer,
@@ -309,4 +314,3 @@ def get_trainer(
callbacks=callbacks,
formatting_func=partial(_formatting_fn, tokenizer=tokenizer),
)
-
diff --git a/src/atgen/utils/validate_and_fill_config.py b/src/atgen/utils/validate_and_fill_config.py
index 4206679..9fd982e 100644
--- a/src/atgen/utils/validate_and_fill_config.py
+++ b/src/atgen/utils/validate_and_fill_config.py
@@ -1,7 +1,10 @@
import os
-from omegaconf import DictConfig, open_dict
+from omegaconf import OmegaConf, DictConfig, open_dict
import logging
+from .data.get_output_column_name_for_phase import get_output_column_name_for_phase
+from .constants import OUTPUT_FIELD_PURPOSE_TRAIN, OUTPUT_FIELD_PURPOSE_TEST
+
log = logging.getLogger(__name__)
@@ -58,6 +61,21 @@ def validate_and_fill_config(config: DictConfig) -> DictConfig:
config.data.setdefault("few_shot", {})
config.data.few_shot.setdefault("count", 0)
config.data.few_shot.setdefault("separator", "\n\n")
+ # Add train and test output column names if not provided
+ if "train_output_column_name" not in config.data:
+ OmegaConf.update(
+ config,
+ "data.train_output_column_name",
+ get_output_column_name_for_phase(config.data.output_column_name, OUTPUT_FIELD_PURPOSE_TRAIN),
+ force_add=True,
+ )
+ if "test_output_column_name" not in config.data:
+ OmegaConf.update(
+ config,
+ "data.test_output_column_name",
+ get_output_column_name_for_phase(config.data.output_column_name, OUTPUT_FIELD_PURPOSE_TEST),
+ force_add=True,
+ )
# Model
validate_field(config, "model")
diff --git a/test/test_benchmark.py b/test/test_benchmark.py
index d5b517c..799cc81 100644
--- a/test/test_benchmark.py
+++ b/test/test_benchmark.py
@@ -17,7 +17,7 @@ def exec_bash(s):
def test_just_works():
exec_result = exec_bash(
- "HYDRA_CONFIG_NAME=test run-al +debug=false +save_model=false"
+ "HYDRA_CONFIG_NAME=test run-al +save_model=false"
)
assert exec_result.returncode == 0
diff --git a/test/test_metrics.py b/test/test_metrics.py
new file mode 100644
index 0000000..53125e2
--- /dev/null
+++ b/test/test_metrics.py
@@ -0,0 +1,785 @@
+import pytest
+import numpy as np
+from unittest.mock import Mock, patch
+import os
+import tempfile
+import shutil
+import warnings
+
+from atgen.metrics import (
+ BaseMetric,
+ MetricConfig,
+ MetricsConfig,
+ MetricsFactory,
+ BleuMetric,
+ RougeMetric,
+ SentBertMetric,
+ ColaMetric,
+ BartScoreMetric,
+ AlignScoreMetric,
+ compute_metrics,
+ compute_metrics_from_config,
+ get_default_config,
+ get_comprehensive_config,
+ get_metric_categories,
+ get_metric_requirements,
+ AVAILABLE_METRICS,
+ get_available_metrics,
+ get_all_possible_metric_keys,
+)
+from atgen.utils.check_required_performance import check_required_performance
+from atgen.utils.check_performance_metrics import check_performance_against_requirements
+from omegaconf import DictConfig
+
+
+class TestMetricConfig:
+ """Test MetricConfig dataclass."""
+
+ def test_default_config(self):
+ """Test default configuration values."""
+ config = MetricConfig()
+ assert config.batch_size == 32
+ assert config.device == "cuda"
+ assert config.cache_dir == "cache"
+ assert config.aggregate is True
+ assert config.threshold == 0.5
+ assert config.provider is None
+ assert config.api_key is None
+
+ def test_custom_config(self):
+ """Test custom configuration values."""
+ config = MetricConfig(
+ batch_size=16,
+ device="cpu",
+ cache_dir="custom_cache",
+ aggregate=False,
+ threshold=0.7,
+ provider="openrouter",
+ api_key="test-key"
+ )
+ assert config.batch_size == 16
+ assert config.device == "cpu"
+ assert config.cache_dir == "custom_cache"
+ assert config.aggregate is False
+ assert config.threshold == 0.7
+ assert config.provider == "openrouter"
+ assert config.api_key == "test-key"
+
+
+class TestMetricsConfig:
+ """Test MetricsConfig dataclass."""
+
+ def test_default_config(self):
+ """Test default configuration values."""
+ config = MetricsConfig()
+ assert config.batch_size == 32
+ assert config.device == "cuda"
+ assert config.cache_dir == "cache"
+ assert config.aggregate is True
+ assert config.metrics == []
+
+ def test_additional_metrics_merge(self):
+ """Test that additional_metrics are merged into metrics."""
+ config = MetricsConfig(
+ metrics=["bleu", "rouge1"],
+ additional_metrics=["sentbert", "cola"]
+ )
+ expected_metrics = ["bleu", "rouge1", "sentbert", "cola"]
+ assert config.metrics == expected_metrics
+
+ def test_duplicate_metrics_removed(self):
+ """Test that duplicate metrics are removed while preserving order."""
+ config = MetricsConfig(
+ metrics=["bleu", "rouge1", "bleu"],
+ additional_metrics=["sentbert", "rouge1"]
+ )
+ expected_metrics = ["bleu", "rouge1", "sentbert"]
+ assert config.metrics == expected_metrics
+
+ def test_deepeval_legacy_params(self):
+ """Test that legacy DeepEval parameters are mapped correctly."""
+ config = MetricsConfig(
+ deepeval_threshold=0.8,
+ deepeval_include_reason=True,
+ deepeval_strict_mode=True,
+ deepeval_async_mode=False,
+ deepeval_verbose_mode=True,
+ deepeval_truths_extraction_limit=10
+ )
+ assert config.threshold == 0.8
+ assert config.include_reason is True
+ assert config.strict_mode is True
+ assert config.async_mode is False
+ assert config.verbose_mode is True
+ assert config.truths_extraction_limit == 10
+
+
+class TestMetricsFactory:
+ """Test MetricsFactory class."""
+
+ def test_get_available_metrics(self):
+ """Test getting available metrics."""
+ metrics = MetricsFactory.get_available_metrics()
+ assert isinstance(metrics, list)
+ assert len(metrics) > 0
+ assert "bleu" in metrics
+ assert "rouge1" in metrics
+ assert "sentbert" in metrics
+ assert "cola" in metrics
+
+ def test_get_all_possible_metric_keys(self):
+ """Test getting all possible metric keys."""
+ keys = MetricsFactory.get_all_possible_metric_keys()
+ assert isinstance(keys, list)
+ assert len(keys) > len(MetricsFactory.get_available_metrics())
+
+ # Should include specific metric keys
+ assert "sentbert_pred_ref" in keys
+ assert "sentbert_pred_src" in keys
+ assert "rouge1" in keys
+ assert "rouge2" in keys
+ assert "rougeL" in keys
+ assert "exact_match" in keys
+ assert "word_length_gen" in keys
+
+ def test_available_metrics_global(self):
+ """Test AVAILABLE_METRICS global variable."""
+ assert isinstance(AVAILABLE_METRICS, list)
+ assert len(AVAILABLE_METRICS) > 20 # Should be comprehensive
+ assert "sentbert_pred_ref" in AVAILABLE_METRICS
+ assert "deepeval_answer_relevance" in AVAILABLE_METRICS
+
+ def test_create_metric_success(self):
+ """Test successful metric creation."""
+ config = MetricConfig(device="cpu")
+ metric = MetricsFactory.create_metric("bleu", config)
+ assert isinstance(metric, BleuMetric)
+ assert metric.config.device == "cpu"
+
+ def test_create_metric_case_insensitive(self):
+ """Test that metric creation is case insensitive."""
+ metric_lower = MetricsFactory.create_metric("bleu")
+ metric_upper = MetricsFactory.create_metric("BLEU")
+ metric_mixed = MetricsFactory.create_metric("Bleu")
+
+ assert type(metric_lower) == type(metric_upper) == type(metric_mixed)
+
+ def test_create_metric_unknown(self):
+ """Test error handling for unknown metrics."""
+ with pytest.raises(ValueError, match="Unknown metric"):
+ MetricsFactory.create_metric("unknown_metric")
+
+ def test_create_metrics_from_config(self):
+ """Test creating multiple metrics from config."""
+ config = MetricsConfig(
+ metrics=["bleu", "rouge1", "sentbert"],
+ device="cpu"
+ )
+ metrics = MetricsFactory.create_metrics(config)
+
+ assert len(metrics) == 3
+ assert "bleu" in metrics
+ assert "rouge1" in metrics
+ assert "sentbert" in metrics
+ assert all(isinstance(m, BaseMetric) for m in metrics.values())
+
+ def test_create_from_config_dict(self):
+ """Test creating metrics from dictionary config."""
+ config_dict = {
+ "metrics": ["bleu", "rouge1"],
+ "device": "cpu",
+ "batch_size": 16
+ }
+ metrics = MetricsFactory.create_from_config(config_dict)
+
+ assert len(metrics) == 2
+ assert "bleu" in metrics
+ assert "rouge1" in metrics
+
+
+class TestPerformanceCheckingIntegration:
+ """Test integration with performance checking system."""
+
+ def test_check_required_performance_valid_metrics(self):
+ """Test check_required_performance with valid metrics."""
+ performance_dict = DictConfig({
+ 'bleu': 0.3,
+ 'rouge1': 0.5,
+ 'sentbert_pred_ref': 0.7,
+ 'deepeval_answer_relevance': 0.8
+ })
+
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ result = check_required_performance(performance_dict)
+
+ # Should have no warnings for valid metrics
+ assert len(w) == 0
+ assert dict(result) == dict(performance_dict)
+
+ def test_check_required_performance_invalid_metrics(self):
+ """Test check_required_performance with invalid metrics."""
+ performance_dict = DictConfig({
+ 'bleu': 0.3,
+ 'nonexistent_metric': 0.5,
+ 'rouge1': 1.5 # Invalid value > 1
+ })
+
+ with warnings.catch_warnings(record=True) as w:
+ warnings.simplefilter("always")
+ result = check_required_performance(performance_dict)
+
+ # Should have warnings for invalid metrics
+ assert len(w) == 2
+ assert dict(result) == {'bleu': 0.3} # Only valid metric remains
+
+ def test_performance_checking_pipeline(self):
+ """Test complete performance checking pipeline."""
+ # Simulate computed metrics
+ metrics = {
+ 'bleu': 0.4,
+ 'rouge1': 0.6,
+ 'sentbert_pred_ref': 0.8,
+ 'word_length_gen': 10.5,
+ 'exact_match': 0.2,
+ 'time_total': 5.2
+ }
+
+ # Set performance requirements
+ required_performance = DictConfig({
+ 'bleu': 0.3,
+ 'rouge1': 0.5,
+ 'sentbert_pred_ref': 0.7
+ })
+
+ # Check requirements
+ checked_requirements = check_required_performance(required_performance)
+
+ # Test performance pipeline
+ is_performance_reached, is_metrics_availability_checked, available_metrics = (
+ check_performance_against_requirements(
+ metrics=metrics,
+ required_performance_dict=checked_requirements,
+ is_metrics_availability_checked=False,
+ available_metrics={}
+ )
+ )
+
+ assert is_performance_reached is True
+ assert is_metrics_availability_checked is True
+ assert len(available_metrics) == 3
+
+
+class TestBasicMetricFunctionality:
+ """Test basic functionality of individual metrics."""
+
+ @pytest.fixture
+ def sample_data(self):
+ """Sample data for testing."""
+ return {
+ "predictions": ["This is a test sentence.", "Another test sentence here."],
+ "references": ["This is a test sentence.", "Another test sentence here."],
+ "original_texts": ["What is this?", "What is that?"]
+ }
+
+ @pytest.fixture
+ def temp_cache_dir(self):
+ """Create temporary cache directory for tests."""
+ temp_dir = tempfile.mkdtemp()
+ yield temp_dir
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
+ def test_bleu_metric(self, sample_data, temp_cache_dir):
+ """Test BLEU metric functionality."""
+ config = MetricConfig(device="cpu", cache_dir=temp_cache_dir)
+ metric = BleuMetric(config)
+
+ results = metric.calculate(
+ sample_data["predictions"],
+ sample_data["references"],
+ sample_data["original_texts"]
+ )
+
+ assert "bleu" in results
+ if isinstance(results["bleu"], np.ndarray):
+ assert all(score == 1.0 for score in results["bleu"]) # Identical strings
+ else:
+ assert results["bleu"] == 1.0
+
+ def test_rouge_metric(self, sample_data, temp_cache_dir):
+ """Test ROUGE metric functionality."""
+ config = MetricConfig(device="cpu", cache_dir=temp_cache_dir)
+ metric = RougeMetric(config)
+
+ results = metric.calculate(
+ sample_data["predictions"],
+ sample_data["references"],
+ sample_data["original_texts"]
+ )
+
+ # Should return multiple ROUGE variants
+ rouge_keys = ["rouge1", "rouge2", "rougeL", "rougeLsum"]
+ found_keys = [key for key in rouge_keys if key in results]
+ assert len(found_keys) > 0
+
+ # Perfect scores for identical strings
+ for key in found_keys:
+ if isinstance(results[key], np.ndarray):
+ assert all(score == 1.0 for score in results[key])
+ else:
+ assert results[key] == 1.0
+
+ @patch('torch.cuda.is_available', return_value=False)
+ def test_sentbert_metric(self, mock_cuda, sample_data, temp_cache_dir):
+ """Test SentBERT metric functionality."""
+ config = MetricConfig(device="cpu", cache_dir=temp_cache_dir)
+ metric = SentBertMetric(config)
+
+ results = metric.calculate(
+ sample_data["predictions"],
+ sample_data["references"],
+ sample_data["original_texts"]
+ )
+
+ # Should return both pred-ref and pred-src similarities
+ assert "sentbert_pred_ref" in results
+ assert "sentbert_pred_src" in results
+
+ # High similarity for identical strings
+ if isinstance(results["sentbert_pred_ref"], np.ndarray):
+ assert all(score > 0.99 for score in results["sentbert_pred_ref"])
+ else:
+ assert results["sentbert_pred_ref"] > 0.99
+
+ @patch('torch.cuda.is_available', return_value=False)
+ def test_cola_metric(self, mock_cuda, sample_data, temp_cache_dir):
+ """Test CoLA metric functionality."""
+ config = MetricConfig(device="cpu", cache_dir=temp_cache_dir)
+ metric = ColaMetric(config)
+
+ results = metric.calculate(
+ sample_data["predictions"],
+ sample_data["references"],
+ sample_data["original_texts"]
+ )
+
+ assert "cola" in results or "grammaticality" in results
+
+ # Should return reasonable grammaticality scores
+ score_key = "cola" if "cola" in results else "grammaticality"
+ if isinstance(results[score_key], np.ndarray):
+ assert all(0 <= score <= 1 for score in results[score_key])
+ else:
+ assert 0 <= results[score_key] <= 1
+
+
+class TestComputeMetrics:
+ """Test the main compute_metrics function."""
+
+ @pytest.fixture
+ def sample_data(self):
+ """Sample data for testing."""
+ return {
+ "predictions": ["This is a test.", "Another test."],
+ "references": ["This is a test.", "Another test."],
+ "original_texts": ["Source text one.", "Source text two."]
+ }
+
+ @pytest.fixture
+ def temp_cache_dir(self):
+ """Create temporary cache directory for tests."""
+ temp_dir = tempfile.mkdtemp()
+ yield temp_dir
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
+ def test_compute_metrics_default_config(self, sample_data, temp_cache_dir):
+ """Test compute_metrics with default configuration."""
+ results = compute_metrics(
+ generated_texts=sample_data["predictions"],
+ reference_texts=sample_data["references"],
+ original_texts=sample_data["original_texts"],
+ cache_dir=temp_cache_dir
+ )
+
+ assert isinstance(results, dict)
+ assert len(results) > 0
+
+ # Should include timing information
+ assert any(key.startswith("time_") for key in results.keys())
+ assert "time_total" in results
+
+ # Should include basic statistics
+ assert "word_length_gen" in results
+ assert "exact_match" in results
+
+ # Exact match should be 1.0 for identical strings
+ assert results["exact_match"] == 1.0
+
+ def test_compute_metrics_comprehensive_config(self, sample_data, temp_cache_dir):
+ """Test compute_metrics with comprehensive configuration."""
+ config = DictConfig({
+ 'metrics': ['bleu', 'rouge1', 'sentbert', 'cola'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir,
+ 'batch_size': 16,
+ 'aggregate': True
+ })
+
+ results = compute_metrics(
+ generated_texts=sample_data["predictions"],
+ reference_texts=sample_data["references"],
+ original_texts=sample_data["original_texts"],
+ config=config
+ )
+
+ assert isinstance(results, dict)
+ assert len(results) > 10 # Should have many metrics
+
+ # Check key metrics are present
+ assert "bleu" in results
+ assert "rouge1" in results
+ assert "sentbert_pred_ref" in results
+ assert any(key in results for key in ["cola", "grammaticality"])
+
+ # Perfect scores for identical strings
+ assert results["bleu"] == 1.0
+ assert results["rouge1"] == 1.0
+ assert results["exact_match"] == 1.0
+
+ def test_compute_metrics_with_deepeval_config(self, sample_data, temp_cache_dir):
+ """Test compute_metrics with DeepEval configuration."""
+ # Skip if no API key available
+ api_key = os.environ.get('OPENROUTER_API_KEY')
+ if not api_key:
+ pytest.skip("No OpenRouter API key available for DeepEval testing")
+
+ config = DictConfig({
+ 'metrics': ['bleu', 'deepeval_answer_relevance'],
+ 'provider': 'openrouter',
+ 'base_url': 'https://openrouter.ai/api/v1',
+ 'api_key': api_key,
+ 'model': 'openai/gpt-4o-mini',
+ 'threshold': 0.5,
+ 'async_mode': False,
+ 'verbose_mode': False,
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir
+ })
+
+ results = compute_metrics(
+ generated_texts=sample_data["predictions"],
+ reference_texts=sample_data["references"],
+ original_texts=sample_data["original_texts"],
+ config=config
+ )
+
+ assert isinstance(results, dict)
+ assert "bleu" in results
+
+ # Check if DeepEval worked
+ deepeval_keys = [k for k in results.keys() if 'deepeval' in k and not k.startswith('time_')]
+ if deepeval_keys:
+ assert len(deepeval_keys) > 0
+ for key in deepeval_keys:
+ assert 0 <= results[key] <= 1
+
+ def test_compute_metrics_no_references(self, sample_data, temp_cache_dir):
+ """Test compute_metrics without references."""
+ config = DictConfig({
+ 'metrics': ['sentbert', 'cola'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir
+ })
+
+ results = compute_metrics(
+ generated_texts=sample_data["predictions"],
+ reference_texts=None,
+ original_texts=sample_data["original_texts"],
+ config=config
+ )
+
+ assert isinstance(results, dict)
+ assert "word_length_gen" in results
+ assert "sentbert_pred_src" in results # Should have pred-src similarity
+
+ def test_compute_metrics_empty_predictions(self):
+ """Test compute_metrics with empty predictions."""
+ with pytest.raises(ValueError, match="predictions cannot be empty"):
+ compute_metrics(generated_texts=[])
+
+ def test_compute_metrics_from_config(self, sample_data, temp_cache_dir):
+ """Test compute_metrics_from_config function."""
+ config_dict = {
+ "metrics": ["bleu", "rouge1"],
+ "device": "cpu",
+ "cache_dir": temp_cache_dir
+ }
+
+ results = compute_metrics_from_config(
+ predictions=sample_data["predictions"],
+ references=sample_data["references"],
+ original_texts=sample_data["original_texts"],
+ config_dict=config_dict
+ )
+
+ assert isinstance(results, dict)
+ assert len(results) > 0
+ assert "bleu" in results
+ assert "rouge1" in results
+
+
+class TestMetricRequirements:
+ """Test metric requirements and categories."""
+
+ def test_get_metric_categories(self):
+ """Test getting metric categories."""
+ categories = get_metric_categories()
+
+ assert isinstance(categories, dict)
+ assert "lexical" in categories
+ assert "semantic" in categories
+ assert "linguistic" in categories
+ assert "llm_based" in categories
+
+ # Check specific metrics in categories
+ assert "bleu" in categories["lexical"]
+ assert "rouge" in categories["lexical"]
+ assert "sentbert" in categories["semantic"]
+ assert "cola" in categories["linguistic"]
+
+ def test_get_metric_requirements(self):
+ """Test getting metric requirements."""
+ requirements = get_metric_requirements()
+
+ assert isinstance(requirements, dict)
+
+ # Test specific requirements
+ assert requirements["bleu"]["requires_references"] is True
+ assert requirements["bleu"]["requires_original_texts"] is False
+
+ assert requirements["sentbert"]["requires_references"] is False
+ assert requirements["sentbert"]["requires_original_texts"] is False
+
+ assert requirements["cola"]["requires_references"] is False
+ assert requirements["cola"]["requires_original_texts"] is False
+
+
+class TestConfigurationFunctions:
+ """Test configuration helper functions."""
+
+ def test_get_default_config(self):
+ """Test get_default_config function."""
+ config = get_default_config()
+
+ assert isinstance(config, MetricsConfig)
+ assert "bleu" in config.metrics
+ assert "rouge1" in config.metrics or any("rouge" in m for m in config.metrics)
+ assert config.batch_size == 32
+ assert config.device == "cuda"
+
+ def test_get_comprehensive_config(self):
+ """Test get_comprehensive_config function."""
+ config = get_comprehensive_config()
+
+ assert isinstance(config, MetricsConfig)
+ assert len(config.metrics) > 2
+ assert "bleu" in config.metrics
+ assert any("rouge" in metric for metric in config.metrics)
+ assert "sentbert" in config.metrics
+ assert "cola" in config.metrics
+
+
+class TestEdgeCases:
+ """Test edge cases and error handling."""
+
+ @pytest.fixture
+ def temp_cache_dir(self):
+ """Create temporary cache directory for tests."""
+ temp_dir = tempfile.mkdtemp()
+ yield temp_dir
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
+ def test_empty_strings(self, temp_cache_dir):
+ """Test metrics with empty strings."""
+ config = DictConfig({
+ 'metrics': ['bleu', 'rouge1'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir
+ })
+
+ results = compute_metrics(
+ generated_texts=["", ""],
+ reference_texts=["", ""],
+ original_texts=["", ""],
+ config=config
+ )
+
+ assert isinstance(results, dict)
+ # Should handle empty strings gracefully
+
+ def test_single_item_lists(self, temp_cache_dir):
+ """Test metrics with single item lists."""
+ config = DictConfig({
+ 'metrics': ['bleu', 'rouge1'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir
+ })
+
+ results = compute_metrics(
+ generated_texts=["Single test sentence."],
+ reference_texts=["Single test sentence."],
+ original_texts=["Single source text."],
+ config=config
+ )
+
+ assert isinstance(results, dict)
+ assert len(results) > 0
+ assert results["exact_match"] == 1.0 # Identical strings
+
+ def test_multiple_references(self, temp_cache_dir):
+ """Test metrics with multiple references."""
+ config = DictConfig({
+ 'metrics': ['bleu', 'rouge1'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir
+ })
+
+ results = compute_metrics(
+ generated_texts=["Test sentence."],
+ reference_texts=[["Test sentence.", "Another reference."]],
+ original_texts=["Source text."],
+ config=config
+ )
+
+ assert isinstance(results, dict)
+ assert len(results) > 0
+
+
+class TestRealWorldScenarios:
+ """Test realistic scenarios that would occur in active learning."""
+
+ @pytest.fixture
+ def temp_cache_dir(self):
+ """Create temporary cache directory for tests."""
+ temp_dir = tempfile.mkdtemp()
+ yield temp_dir
+ shutil.rmtree(temp_dir, ignore_errors=True)
+
+ def test_active_learning_integration(self, temp_cache_dir):
+ """Test complete active learning integration scenario."""
+ # Realistic generated texts
+ generated_texts = [
+ "Machine learning algorithms learn patterns from training data to make predictions.",
+ "Deep learning uses neural networks with multiple layers to process information.",
+ "Natural language processing helps computers understand human language."
+ ]
+
+ reference_texts = [
+ "ML algorithms learn from data to make predictions on new examples.",
+ "Deep learning employs multi-layered neural networks for complex data processing.",
+ "NLP enables computers to work with human language effectively."
+ ]
+
+ original_texts = [
+ "How do machine learning algorithms work?",
+ "What is deep learning?",
+ "Explain natural language processing."
+ ]
+
+ # Comprehensive config similar to active learning setup
+ config = DictConfig({
+ 'metrics': ['bleu', 'rouge1', 'rouge2', 'rougeL', 'sentbert'],
+ 'additional_metrics': ['cola'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir,
+ 'batch_size': 16,
+ 'aggregate': True
+ })
+
+ # Compute metrics
+ results = compute_metrics(
+ generated_texts=generated_texts,
+ reference_texts=reference_texts,
+ original_texts=original_texts,
+ config=config
+ )
+
+ # Verify comprehensive results
+ assert isinstance(results, dict)
+ assert len(results) > 10
+
+ # Check key performance metrics
+ performance_keys = ['bleu', 'rouge1', 'rouge2', 'rougeL', 'sentbert_pred_ref']
+ for key in performance_keys:
+ if key in results:
+ assert isinstance(results[key], (int, float))
+ assert 0 <= results[key] <= 1
+
+ # Test performance checking integration
+ required_performance = DictConfig({
+ 'bleu': 0.15,
+ 'rouge1': 0.35,
+ 'sentbert_pred_ref': 0.55
+ })
+
+ checked_requirements = check_required_performance(required_performance)
+ assert len(checked_requirements) == 3 # All should be valid
+
+ # Test complete pipeline
+ is_performance_reached, is_metrics_availability_checked, available_metrics = (
+ check_performance_against_requirements(
+ metrics=results,
+ required_performance_dict=checked_requirements,
+ is_metrics_availability_checked=False,
+ available_metrics={}
+ )
+ )
+
+ assert is_metrics_availability_checked is True
+ assert len(available_metrics) == 3
+
+ def test_high_quality_generation_scenario(self, temp_cache_dir):
+ """Test scenario with high-quality generated text."""
+ # High-quality generated texts (should score well)
+ generated_texts = [
+ "The capital of France is Paris, which is also its largest city.",
+ "Artificial intelligence is transforming various industries worldwide."
+ ]
+
+ reference_texts = [
+ "Paris is the capital and largest city of France.",
+ "AI is revolutionizing industries across the globe."
+ ]
+
+ original_texts = [
+ "What is the capital of France?",
+ "How is AI impacting industries?"
+ ]
+
+ config = DictConfig({
+ 'metrics': ['bleu', 'rouge1', 'sentbert', 'cola'],
+ 'device': 'cpu',
+ 'cache_dir': temp_cache_dir,
+ 'aggregate': True
+ })
+
+ results = compute_metrics(
+ generated_texts=generated_texts,
+ reference_texts=reference_texts,
+ original_texts=original_texts,
+ config=config
+ )
+
+ # Should get reasonable scores for good quality text
+ assert results['bleu'] > 0.1 # Some lexical overlap
+ assert results['rouge1'] > 0.2 # Some word overlap
+ assert results['sentbert_pred_ref'] > 0.5 # Good semantic similarity
+
+ # CoLA should give high grammaticality scores
+ cola_key = 'cola' if 'cola' in results else 'grammaticality'
+ if cola_key in results:
+ assert results[cola_key] > 0.7 # Well-formed sentences
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
\ No newline at end of file