valtron-core 0.1.0__py3-none-any.whl

valtron_core/__init__.py

@@ -0,0 +1,3 @@
+"""Valtron Core: LLM call optimization across various providers."""
+
+__version__ = "0.1.0"

valtron_core/__main__.py

@@ -0,0 +1,16 @@
+"""Main entry point for the evaltron_core package."""
+
+import typer
+
+app = typer.Typer()
+
+
+@app.command()
+def main() -> None:
+    """Main entry point for evaltron_core."""
+    typer.echo("Valtron Core - LLM call optimization")
+    typer.echo("Project initialized successfully!")
+
+
+if __name__ == "__main__":
+    app()

valtron_core/attachments.py

@@ -0,0 +1,32 @@
+"""Shared attachment MIME detection utilities."""
+
+from pathlib import Path
+
+_EXT_MIME: dict[str, str] = {
+    ".pdf": "application/pdf",
+    ".png": "image/png",
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".gif": "image/gif",
+    ".webp": "image/webp",
+}
+
+_MAGIC: list[tuple[bytes, str]] = [
+    (b"%PDF-", "application/pdf"),
+    (b"\x89PNG", "image/png"),
+    (b"\xff\xd8\xff", "image/jpeg"),
+    (b"GIF8", "image/gif"),
+    (b"RIFF", "image/webp"),
+]
+
+
+def detect_mime_hint(s: str) -> str:
+    """
+    Detect MIME type from a data URI header or file/URL extension without any I/O.
+    Returns an empty string if the type cannot be determined.
+    """
+    if s.startswith("data:"):
+        header = s.split(",")[0]
+        return header.split(":")[1].split(";")[0]
+    suffix = Path(s.split("?")[0]).suffix.lower()
+    return _EXT_MIME.get(suffix, "")

valtron_core/bert_evaluator.py

@@ -0,0 +1,227 @@
+"""BERT model evaluation integrated with the evaluation framework.
+
+.. deprecated::
+    Use :class:`evaltron_core.transformer_classifier.TransformerClassifier` for
+    inference and the recipe layer for evaluation. ``BERTEvaluator`` will be
+    removed in a future release.
+"""
+
+import warnings
+
+warnings.warn(
+    "evaltron_core.bert_evaluator.BERTEvaluator is deprecated. "
+    "Use evaltron_core.transformer_classifier.TransformerClassifier with "
+    "ModelEval (type='transformer') instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+import time
+from typing import Any
+
+import structlog
+
+from valtron_core.bert_trainer import BERTTrainer
+from valtron_core.models import Document, EvaluationInput, EvaluationResult, Label, PredictionResult
+
+logger = structlog.get_logger()
+
+
+class BERTEvaluator:
+    """Evaluate BERT models using the same framework as LLM evaluations."""
+
+    def __init__(self, trainer: BERTTrainer) -> None:
+        """
+        Initialize BERT evaluator.
+
+        Args:
+            trainer: Trained BERTTrainer instance
+        """
+        self.trainer = trainer
+
+    async def evaluate_single(
+        self,
+        document: Document,
+        label: Label,
+    ) -> PredictionResult:
+        """
+        Evaluate a single document.
+
+        Args:
+            document: Document to evaluate
+            label: Expected label
+
+        Returns:
+            PredictionResult
+        """
+        start_time = time.time()
+
+        try:
+            # Get prediction
+            predicted_value = self.trainer.predict_single(document.content)
+
+            end_time = time.time()
+            response_time = end_time - start_time
+
+            # Compare with expected
+            is_correct = predicted_value.strip().lower() == label.value.strip().lower()
+
+            logger.info(
+                "bert_evaluation_single",
+                document_id=document.id,
+                predicted=predicted_value,
+                expected=label.value,
+                correct=is_correct,
+                time=response_time,
+            )
+
+            return PredictionResult(
+                document_id=document.id,
+                predicted_value=predicted_value,
+                expected_value=label.value,
+                is_correct=is_correct,
+                response_time=response_time,
+                original_cost=0.0,
+                cost=0.0,  # BERT inference is free
+                model="bert-local",
+                metadata={"content": document.content},
+            )
+
+        except Exception as e:
+            end_time = time.time()
+            response_time = end_time - start_time
+
+            logger.error(
+                "bert_evaluation_error",
+                document_id=document.id,
+                error=str(e),
+                time=response_time,
+            )
+
+            return PredictionResult(
+                document_id=document.id,
+                predicted_value=f"ERROR: {str(e)}",
+                expected_value=label.value,
+                is_correct=False,
+                response_time=response_time,
+                original_cost=0.0,
+                cost=0.0,
+                model="bert-local",
+                metadata={"content": document.content, "error": str(e)},
+            )
+
+    async def evaluate(
+        self,
+        eval_input: EvaluationInput,
+    ) -> EvaluationResult:
+        """
+        Evaluate all documents.
+
+        Args:
+            eval_input: Evaluation input configuration
+
+        Returns:
+            EvaluationResult with all predictions and metrics
+        """
+        import uuid
+        from datetime import datetime
+
+        run_id = str(uuid.uuid4())
+        result = EvaluationResult(
+            run_id=run_id,
+            started_at=datetime.now(),
+            prompt_template="BERT local inference (no prompt)",
+            model="bert-local",
+            status="running",
+        )
+
+        # Create label lookup
+        label_map = {label.document_id: label for label in eval_input.labels}
+
+        logger.info(
+            "bert_evaluation_started",
+            run_id=run_id,
+            total_documents=len(eval_input.documents),
+        )
+
+        try:
+            # Evaluate all documents
+            for doc in eval_input.documents:
+                if doc.id not in label_map:
+                    logger.warning("missing_label", document_id=doc.id)
+                    continue
+
+                prediction = await self.evaluate_single(
+                    document=doc,
+                    label=label_map[doc.id],
+                )
+
+                result.predictions.append(prediction)
+
+            # Compute metrics
+            result.compute_metrics()
+            result.completed_at = datetime.now()
+            result.status = "completed"
+
+            logger.info(
+                "bert_evaluation_completed",
+                run_id=run_id,
+                accuracy=result.metrics.accuracy if result.metrics else 0,
+                total_time=result.metrics.total_time if result.metrics else 0,
+            )
+
+        except Exception as e:
+            result.status = "failed"
+            result.error = str(e)
+            result.completed_at = datetime.now()
+
+            logger.error("bert_evaluation_failed", run_id=run_id, error=str(e))
+
+        return result
+
+
+def create_bert_model_for_comparison(
+    documents: list[Document],
+    labels: list[Label],
+    model_name: str = "bert-base-uncased",
+    output_dir: str = "./bert_models",
+    num_epochs: int = 3,
+    batch_size: int = 8,
+) -> BERTTrainer:
+    """
+    Train a BERT model for comparison with LLMs.
+
+    Args:
+        documents: Training documents
+        labels: Training labels
+        model_name: Pretrained BERT model to use
+        output_dir: Directory to save model
+        num_epochs: Number of training epochs
+        batch_size: Training batch size
+
+    Returns:
+        Trained BERTTrainer instance
+    """
+    logger.info("creating_bert_model", model_name=model_name, num_docs=len(documents))
+
+    # Initialize trainer
+    trainer = BERTTrainer(model_name=model_name, output_dir=output_dir)
+
+    # Prepare data
+    train_dataset, test_dataset = trainer.prepare_data(documents, labels)
+
+    # Train
+    results = trainer.train(
+        train_dataset=train_dataset,
+        test_dataset=test_dataset,
+        num_epochs=num_epochs,
+        batch_size=batch_size,
+    )
+
+    logger.info(
+        "bert_model_created",
+        accuracy=results["eval_accuracy"],
+        model_dir=results["model_dir"],
+    )
+
+    return trainer

valtron_core/bert_trainer.py

@@ -0,0 +1,332 @@
+"""BERT model training for classification tasks.
+
+.. deprecated::
+    Use :class:`evaltron_core.transformer_classifier.TransformerClassifier` instead.
+    ``BERTTrainer`` will be removed in a future release.
+"""
+
+import warnings
+
+warnings.warn(
+    "evaltron_core.bert_trainer.BERTTrainer is deprecated. "
+    "Use evaltron_core.transformer_classifier.TransformerClassifier instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+import json
+from pathlib import Path
+from typing import Any
+
+import structlog
+import torch
+from datasets import Dataset
+from sklearn.metrics import accuracy_score, classification_report
+from sklearn.model_selection import train_test_split
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+)
+
+from valtron_core.models import Document, Label
+
+logger = structlog.get_logger()
+
+
+class BERTTrainer:
+    """Train and manage BERT-based classification models."""
+
+    def __init__(
+        self,
+        model_name: str = "bert-base-uncased",
+        output_dir: str | Path = "./bert_models",
+    ) -> None:
+        """
+        Initialize BERT trainer.
+
+        Args:
+            model_name: Pretrained model name from HuggingFace
+            output_dir: Directory to save trained models
+        """
+        self.model_name = model_name
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        self.tokenizer: Any = None
+        self.model: Any = None
+        self.label_to_id: dict[str, int] = {}
+        self.id_to_label: dict[int, str] = {}
+
+    def prepare_data(
+        self,
+        documents: list[Document],
+        labels: list[Label],
+        test_size: float = 0.2,
+        random_state: int = 42,
+    ) -> tuple[Dataset, Dataset]:
+        """
+        Prepare data for training.
+
+        Args:
+            documents: List of documents
+            labels: List of labels
+            test_size: Fraction of data to use for testing
+            random_state: Random seed for reproducibility
+
+        Returns:
+            Tuple of (train_dataset, test_dataset)
+        """
+        # Create label mapping
+        label_map = {label.document_id: label.value for label in labels}
+
+        # Extract texts and labels
+        texts = []
+        label_values = []
+
+        for doc in documents:
+            if doc.id in label_map:
+                texts.append(doc.content)
+                label_values.append(label_map[doc.id])
+
+        # Create unique label mappings
+        unique_labels = sorted(set(label_values))
+        self.label_to_id = {label: idx for idx, label in enumerate(unique_labels)}
+        self.id_to_label = {idx: label for label, idx in self.label_to_id.items()}
+
+        # Convert labels to integers
+        label_ids = [self.label_to_id[label] for label in label_values]
+
+        # Split data
+        train_texts, test_texts, train_labels, test_labels = train_test_split(
+            texts, label_ids, test_size=test_size, random_state=random_state, stratify=label_ids
+        )
+
+        # Create datasets
+        train_dataset = Dataset.from_dict({"text": train_texts, "label": train_labels})
+        test_dataset = Dataset.from_dict({"text": test_texts, "label": test_labels})
+
+        logger.info(
+            "data_prepared",
+            num_train=len(train_dataset),
+            num_test=len(test_dataset),
+            num_labels=len(unique_labels),
+        )
+
+        return train_dataset, test_dataset
+
+    def _tokenize_function(self, examples: dict[str, Any]) -> dict[str, Any]:
+        """Tokenize examples."""
+        return self.tokenizer(examples["text"], padding="max_length", truncation=True)
+
+    def train(
+        self,
+        train_dataset: Dataset,
+        test_dataset: Dataset,
+        num_epochs: int = 3,
+        batch_size: int = 8,
+        learning_rate: float = 2e-5,
+        warmup_steps: int = 500,
+        weight_decay: float = 0.01,
+        save_steps: int = 500,
+        eval_steps: int = 500,
+    ) -> dict[str, Any]:
+        """
+        Train BERT model.
+
+        Args:
+            train_dataset: Training dataset
+            test_dataset: Test dataset
+            num_epochs: Number of training epochs
+            batch_size: Batch size for training
+            learning_rate: Learning rate
+            warmup_steps: Number of warmup steps
+            weight_decay: Weight decay for optimizer
+            save_steps: Save checkpoint every N steps
+            eval_steps: Evaluate every N steps
+
+        Returns:
+            Training metrics and results
+        """
+        # Initialize tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            self.model_name,
+            num_labels=len(self.label_to_id),
+            id2label=self.id_to_label,
+            label2id=self.label_to_id,
+        )
+
+        # Tokenize datasets
+        train_dataset = train_dataset.map(self._tokenize_function, batched=True)
+        test_dataset = test_dataset.map(self._tokenize_function, batched=True)
+
+        # Set format for PyTorch
+        train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+        test_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir=str(self.output_dir),
+            num_train_epochs=num_epochs,
+            per_device_train_batch_size=batch_size,
+            per_device_eval_batch_size=batch_size,
+            warmup_steps=warmup_steps,
+            weight_decay=weight_decay,
+            learning_rate=learning_rate,
+            logging_dir=str(self.output_dir / "logs"),
+            logging_steps=100,
+            eval_strategy="steps",
+            eval_steps=eval_steps,
+            save_steps=save_steps,
+            load_best_model_at_end=True,
+            metric_for_best_model="accuracy",
+        )
+
+        # Compute metrics function
+        def compute_metrics(eval_pred: Any) -> dict[str, float]:
+            predictions, labels = eval_pred
+            predictions = predictions.argmax(-1)
+            accuracy = accuracy_score(labels, predictions)
+            return {"accuracy": accuracy}
+
+        # Create trainer
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=test_dataset,
+            compute_metrics=compute_metrics,
+        )
+
+        # Train
+        logger.info("training_started", epochs=num_epochs, batch_size=batch_size)
+        train_result = trainer.train()
+
+        # Evaluate
+        eval_result = trainer.evaluate()
+
+        logger.info("training_completed", train_loss=train_result.training_loss, **eval_result)
+
+        # Save model
+        final_model_dir = self.output_dir / "final_model"
+        trainer.save_model(str(final_model_dir))
+        self.tokenizer.save_pretrained(str(final_model_dir))
+
+        # Save label mappings
+        label_map_path = final_model_dir / "label_mapping.json"
+        with open(label_map_path, "w") as f:
+            json.dump(
+                {"label_to_id": self.label_to_id, "id_to_label": self.id_to_label},
+                f,
+                indent=2,
+            )
+
+        return {
+            "train_loss": train_result.training_loss,
+            "eval_accuracy": eval_result["eval_accuracy"],
+            "eval_loss": eval_result["eval_loss"],
+            "model_dir": str(final_model_dir),
+        }
+
+    def load_model(self, model_dir: str | Path) -> None:
+        """
+        Load a trained model.
+
+        Args:
+            model_dir: Directory containing saved model
+        """
+        model_dir = Path(model_dir)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
+        self.model = AutoModelForSequenceClassification.from_pretrained(str(model_dir))
+
+        # Load label mappings
+        label_map_path = model_dir / "label_mapping.json"
+        if label_map_path.exists():
+            with open(label_map_path, "r") as f:
+                mappings = json.load(f)
+                self.label_to_id = mappings["label_to_id"]
+                # Convert string keys back to integers for id_to_label
+                self.id_to_label = {int(k): v for k, v in mappings["id_to_label"].items()}
+
+        logger.info("model_loaded", model_dir=str(model_dir))
+
+    def predict(self, texts: list[str]) -> list[str]:
+        """
+        Make predictions on new texts.
+
+        Args:
+            texts: List of texts to classify
+
+        Returns:
+            List of predicted labels
+        """
+        if self.model is None or self.tokenizer is None:
+            raise ValueError("Model not loaded. Train or load a model first.")
+
+        # Tokenize
+        inputs = self.tokenizer(
+            texts,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        # Predict
+        self.model.eval()
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            predictions = outputs.logits.argmax(-1).tolist()
+
+        # Convert to labels
+        predicted_labels = [self.id_to_label[pred] for pred in predictions]
+
+        return predicted_labels
+
+    def predict_single(self, text: str) -> str:
+        """
+        Predict label for a single text.
+
+        Args:
+            text: Text to classify
+
+        Returns:
+            Predicted label
+        """
+        return self.predict([text])[0]
+
+    def evaluate_on_documents(
+        self,
+        documents: list[Document],
+        labels: list[Label],
+    ) -> dict[str, Any]:
+        """
+        Evaluate model on a set of documents.
+
+        Args:
+            documents: List of documents
+            labels: List of expected labels
+
+        Returns:
+            Evaluation metrics
+        """
+        # Create label mapping
+        label_map = {label.document_id: label.value for label in labels}
+
+        # Get predictions
+        texts = [doc.content for doc in documents if doc.id in label_map]
+        expected_labels = [label_map[doc.id] for doc in documents if doc.id in label_map]
+
+        predicted_labels = self.predict(texts)
+
+        # Calculate metrics
+        accuracy = accuracy_score(expected_labels, predicted_labels)
+        report = classification_report(expected_labels, predicted_labels, output_dict=True)
+
+        return {
+            "accuracy": accuracy,
+            "classification_report": report,
+            "num_samples": len(texts),
+        }