Spaces:

compendious
/

precis

Build error

App Files Files Community

compendious commited on Jan 23

Commit

b813321

1 Parent(s): 9e4cab2

Changes

Browse files

Files changed (15) hide show

.gitignore +1 -1
README.md +43 -2
app.py +0 -24
backend/app.py +103 -0
backend/initiate.py +0 -13
requirements.txt +7 -2
scripts/evaluate.py +69 -0
scripts/train.py +91 -0
src/__init__.py +16 -0
src/config.py +67 -0
src/model.py +92 -0
src/tuning/__init__.py +15 -0
src/tuning/data.py +70 -0
src/tuning/lora.py +53 -0
src/tuning/trainer.py +100 -0

.gitignore CHANGED Viewed

	@@ -1,2 +1,2 @@
1	cache
2	- *.~~ipyn~~


1	cache
2	+ *.ipynb

README.md CHANGED Viewed

@@ -2,8 +2,49 @@
 A system for compressing long-form content into clear, structured summaries.
-Précis is designed for articles, papers, and video transcripts. The goal is to be able to extract the meaningful content rather than paraphrase the main ideas.
 ## Model
-The model used is Qwen-2.5-7B-Instruct with 5-bit quantization for efficiency. It's functional for specifically fine-tuning to fit a schema.

 A system for compressing long-form content into clear, structured summaries.
+Précis is designed for articles, papers, and video transcripts. The goal is to extract meaningful content rather than paraphrase main ideas.
 ## Model
+Qwen-2.5-7B-Instruct with 4-bit quantization (BitsAndBytes NF4) for efficiency. Fine-tuned using LoRA for summarization.
+## Installation
+```bash
+pip install -r requirements.txt
+```
+## Usage
+### Training (with dummy data)
+```bash
+# Dry run to validate pipeline
+python scripts/train.py --dry-run
+# Full training
+python scripts/train.py --epochs 3 --batch-size 4
+```
+### Evaluation
+```bash
+python scripts/evaluate.py --checkpoint ./outputs
+```
+## API
+### Running the API
+```bash
+python app.py
+# or
+uvicorn app:app --reload
+```
+### Endpoints
+- `GET /` — API documentation page
+- `GET /health` — Health check
+- `GET /status` — Service status and model info
+- `POST /summarize` — Summarize content from URL (currently returns dummy data)

app.py DELETED Viewed

@@ -1,24 +0,0 @@
-from fastapi import FastAPI
-from fastapi.responses import HTMLResponse
-app = FastAPI(title="Précis — MVP")
-@app.get("/", response_class=HTMLResponse)
-async def root():
-    return """
-    <html>
-        <head>
-            <title>Précis — MVP</title>
-        </head>
-        <body>
-            <h1>Précis — MVP</h1>
-            <p>Welcome to Précis</p>
-        </body>
-    </html>
-    """
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)

backend/app.py ADDED Viewed

	@@ -0,0 +1,103 @@

+"""FastAPI backend for Précis."""
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import HTMLResponse
+from pydantic import BaseModel
+from typing import Optional
+app = FastAPI(
+    title="Précis API",
+    description="Content summarization API",
+    version="0.1.0"
+)
+class SummarizeRequest(BaseModel):
+    """Request model for summarization."""
+    url: str
+    max_length: Optional[int] = 512
+class SummarizeResponse(BaseModel):
+    """Response model for summarization."""
+    url: str
+    summary: str
+    success: bool
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    """Root endpoint with basic info."""
+    return """
+    <!DOCTYPE html>
+    <html>
+        <head>
+            <title>Précis API</title>
+            <style>
+                body { font-family: system-ui; max-width: 800px; margin: 50px auto; padding: 20px; }
+                h1 { color: #333; }
+                code { background: #f4f4f4; padding: 2px 6px; border-radius: 3px; }
+            </style>
+        </head>
+        <body>
+            <h1>Précis API</h1>
+            <p>Content summarization service</p>
+            <h2>Endpoints</h2>
+            <ul>
+                <li><code>POST /summarize</code> - Summarize content from URL</li>
+                <li><code>GET /health</code> - Health check</li>
+                <li><code>GET /status</code> - Service status</li>
+                <li><code>GET /docs</code> - API documentation</li>
+            </ul>
+        </body>
+    </html>
+    """
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    return {"status": "healthy", "service": "precis"}
+@app.get("/status")
+async def status():
+    """Service status endpoint."""
+    return {
+        "service": "Précis API",
+        "version": "0.1.0",
+        "model": "Qwen/Qwen2.5-7B-Instruct",
+        "model_loaded": False,  # TODO: Track actual model state
+        "endpoints": ["/", "/health", "/status", "/summarize"]
+    }
+@app.post("/summarize", response_model=SummarizeResponse)
+async def summarize(request: SummarizeRequest):
+    """
+    Summarize content from a URL.
+    Currently returns dummy data. Will be implemented with actual model.
+    """
+    # TODO: Implement actual summarization
+    # 1. Fetch content from URL
+    # 2. Parse text (YouTube transcript or article)
+    # 3. Run through model
+    # 4. Return summary
+    dummy_summary = (
+        f"This is a placeholder summary for content at {request.url}. "
+        "The actual summarization model will be integrated in the next phase. "
+        "This summary respects the max_length parameter of {request.max_length} tokens."
+    )
+    return SummarizeResponse(
+        url=request.url,
+        summary=dummy_summary,
+        success=True
+    )
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

backend/initiate.py DELETED Viewed

@@ -1,13 +0,0 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-MODEL = "Qwen/Qwen2.5-7B-Instruct.gguf.q5_0"
-tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL,
-    device_map="auto",
-    load_in_4bit=True,
-    torch_dtype="auto",
-    trust_remote_code=True
-)

requirements.txt CHANGED Viewed

@@ -1,7 +1,12 @@
-pytorch
 transformers
 accelerate
 bitsandbytes
-summarizer
 sentencepiece
 fastapi

+# Core ML
+torch
 transformers
 accelerate
 bitsandbytes
+peft
+datasets
 sentencepiece
+# API
 fastapi
+uvicorn

scripts/evaluate.py ADDED Viewed

	@@ -0,0 +1,69 @@

+#!/usr/bin/env python3
+"""CLI evaluation script for Précis."""
+import argparse
+import logging
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from src.config import ModelConfig, DataConfig
+from src.model import load_tokenizer
+from src.tuning.data import create_dummy_data
+from transformers import AutoModelForCausalLM
+from peft import PeftModel
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+def parse_args():
+    parser = argparse.ArgumentParser(description="Evaluate Précis model")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint")
+    parser.add_argument("--num-samples", type=int, default=5, help="Number of samples to evaluate")
+    parser.add_argument("--max-new-tokens", type=int, default=256, help="Max tokens to generate")
+    return parser.parse_args()
+def main():
+    args = parse_args()
+    config = ModelConfig()
+    data_config = DataConfig()
+    logger.info(f"Loading checkpoint from {args.checkpoint}")
+    tokenizer = load_tokenizer(config)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.checkpoint,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    # Generate on dummy samples
+    samples = create_dummy_data(args.num_samples)
+    for i, sample in enumerate(samples):
+        prompt = data_config.format_prompt(sample["text"])
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=args.max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        summary = generated[len(prompt):]
+        logger.info(f"\n=== Sample {i+1} ===")
+        logger.info(f"Input: {sample['text'][:100]}...")
+        logger.info(f"Generated: {summary}")
+        logger.info(f"Reference: {sample['summary']}")
+if __name__ == "__main__":
+    main()

scripts/train.py ADDED Viewed

	@@ -0,0 +1,91 @@

+#!/usr/bin/env python3
+"""CLI training script for Précis."""
+import argparse
+import logging
+import sys
+from pathlib import Path
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from src.config import ModelConfig, TrainingConfig, DataConfig
+from src.model import load_model, load_tokenizer, prepare_for_training
+from src.tuning.lora import apply_lora
+from src.tuning.data import create_dummy_data, prepare_dataset
+from src.tuning.trainer import PrecisTrainer
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train Précis summarization model")
+    parser.add_argument("--model-id", type=str, default=None, help="HuggingFace model ID")
+    parser.add_argument("--output-dir", type=str, default="./outputs", help="Output directory")
+    parser.add_argument("--epochs", type=int, default=3, help="Number of training epochs")
+    parser.add_argument("--batch-size", type=int, default=4, help="Batch size")
+    parser.add_argument("--learning-rate", type=float, default=2e-4, help="Learning rate")
+    parser.add_argument("--lora-r", type=int, default=16, help="LoRA rank")
+    parser.add_argument("--dry-run", action="store_true", help="Validate pipeline without training")
+    parser.add_argument("--dummy-samples", type=int, default=100, help="Number of dummy samples")
+    return parser.parse_args()
+def main():
+    args = parse_args()
+    # Build configs
+    model_config = ModelConfig()
+    if args.model_id:
+        model_config.model_id = args.model_id
+    training_config = TrainingConfig(
+        output_dir=args.output_dir,
+        num_epochs=args.epochs,
+        batch_size=args.batch_size,
+        learning_rate=args.learning_rate,
+        lora_r=args.lora_r,
+    )
+    data_config = DataConfig()
+    if args.dry_run:
+        logger.info("=== DRY RUN MODE ===")
+        logger.info(f"Model: {model_config.model_id}")
+        logger.info(f"Output: {training_config.output_dir}")
+        logger.info(f"Epochs: {training_config.num_epochs}, Batch: {training_config.batch_size}")
+        logger.info(f"LoRA r: {training_config.lora_r}, alpha: {training_config.lora_alpha}")
+        # Test data pipeline only
+        dummy_data = create_dummy_data(5)
+        logger.info(f"Dummy data sample: {dummy_data[0]}")
+        logger.info("Dry run complete. Pipeline validated.")
+        return
+    # Load model and tokenizer
+    logger.info("Loading model and tokenizer...")
+    tokenizer = load_tokenizer(model_config)
+    model = load_model(model_config)
+    model = prepare_for_training(model)
+    model = apply_lora(model, training_config)
+    # Prepare data
+    logger.info("Preparing training data...")
+    train_data = create_dummy_data(args.dummy_samples)
+    train_dataset = prepare_dataset(train_data, tokenizer, data_config)
+    # Train
+    trainer = PrecisTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=train_dataset,
+        config=training_config,
+    )
+    trainer.train()
+    trainer.save()
+    logger.info("Training complete!")
+if __name__ == "__main__":
+    main()

src/__init__.py ADDED Viewed

	@@ -0,0 +1,16 @@

+"""
+Précis — Model loading, configuration, and fine-tuning utilities.
+"""
+from src.config import ModelConfig, TrainingConfig, DataConfig
+from src.model import load_model, load_tokenizer, prepare_for_training
+__version__ = "0.1.0"
+__all__ = [
+    "ModelConfig",
+    "TrainingConfig",
+    "DataConfig",
+    "load_model",
+    "load_tokenizer",
+    "prepare_for_training",
+]

src/config.py ADDED Viewed

	@@ -0,0 +1,67 @@

+"""Configuration management for Précis."""
+from dataclasses import dataclass, field
+from typing import Optional, List
+@dataclass
+class ModelConfig:
+    """Configuration for model loading and quantization."""
+    model_id: str = "Qwen/Qwen2.5-7B-Instruct"
+    load_in_4bit: bool = True
+    load_in_8bit: bool = False
+    bnb_4bit_compute_dtype: str = "float16"
+    bnb_4bit_quant_type: str = "nf4"
+    bnb_4bit_use_double_quant: bool = True
+    device_map: str = "auto"
+    trust_remote_code: bool = True
+    cache_dir: Optional[str] = None
+    def __post_init__(self):
+        if self.load_in_4bit and self.load_in_8bit:
+            raise ValueError("Cannot enable both 4-bit and 8-bit quantization")
+@dataclass
+class TrainingConfig:
+    """Configuration for LoRA fine-tuning."""
+    lora_r: int = 16
+    lora_alpha: int = 32
+    lora_dropout: float = 0.05
+    lora_target_modules: List[str] = field(
+        default_factory=lambda: ["q_proj", "k_proj", "v_proj", "o_proj"]
+    )
+    learning_rate: float = 2e-4
+    batch_size: int = 4
+    gradient_accumulation_steps: int = 4
+    num_epochs: int = 3
+    warmup_ratio: float = 0.03
+    weight_decay: float = 0.01
+    max_grad_norm: float = 1.0
+    max_seq_length: int = 2048
+    optim: str = "paged_adamw_32bit"
+    save_steps: int = 100
+    logging_steps: int = 10
+    eval_steps: int = 100
+    output_dir: str = "./outputs"
+    seed: int = 42
+@dataclass
+class DataConfig:
+    """Configuration for dataset loading and preprocessing."""
+    train_file: Optional[str] = None
+    eval_file: Optional[str] = None
+    input_column: str = "text"
+    target_column: str = "summary"
+    max_input_length: int = 1536
+    max_target_length: int = 512
+    train_split: float = 0.9
+    prompt_template: str = (
+        "Summarize the following document:\n\n"
+        "### Document:\n{input}\n\n"
+        "### Summary:\n"
+    )
+    def format_prompt(self, text: str) -> str:
+        return self.prompt_template.format(input=text)

src/model.py ADDED Viewed

	@@ -0,0 +1,92 @@

+"""Model loading utilities for Précis."""
+import logging
+from typing import Optional, Tuple
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+)
+from src.config import ModelConfig
+logger = logging.getLogger(__name__)
+def get_quantization_config(config: ModelConfig) -> Optional[BitsAndBytesConfig]:
+    """Create BitsAndBytes quantization configuration."""
+    if config.load_in_4bit:
+        compute_dtype = getattr(torch, config.bnb_4bit_compute_dtype)
+        return BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=compute_dtype,
+            bnb_4bit_quant_type=config.bnb_4bit_quant_type,
+            bnb_4bit_use_double_quant=config.bnb_4bit_use_double_quant,
+        )
+    elif config.load_in_8bit:
+        return BitsAndBytesConfig(load_in_8bit=True)
+    return None
+def load_tokenizer(config: Optional[ModelConfig] = None) -> PreTrainedTokenizer:
+    """Load and configure the tokenizer."""
+    if config is None:
+        config = ModelConfig()
+    logger.info(f"Loading tokenizer: {config.model_id}")
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.model_id,
+        trust_remote_code=config.trust_remote_code,
+        cache_dir=config.cache_dir,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    tokenizer.padding_side = "right"
+    return tokenizer
+def load_model(config: Optional[ModelConfig] = None) -> PreTrainedModel:
+    """Load the base model with optional quantization."""
+    if config is None:
+        config = ModelConfig()
+    logger.info(f"Loading model: {config.model_id}")
+    quantization_config = get_quantization_config(config)
+    model = AutoModelForCausalLM.from_pretrained(
+        config.model_id,
+        quantization_config=quantization_config,
+        device_map=config.device_map,
+        trust_remote_code=config.trust_remote_code,
+        cache_dir=config.cache_dir,
+        torch_dtype=torch.float16 if quantization_config else "auto",
+    )
+    logger.info(f"Model loaded. Parameters: {model.num_parameters():,}")
+    return model
+def prepare_for_training(model: PreTrainedModel, gradient_checkpointing: bool = True) -> PreTrainedModel:
+    """Prepare model for training with gradient checkpointing and k-bit setup."""
+    if gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+    if getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False):
+        from peft import prepare_model_for_kbit_training
+        model = prepare_model_for_kbit_training(model)
+    return model
+def load_model_and_tokenizer(config: Optional[ModelConfig] = None) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
+    """Load both model and tokenizer."""
+    if config is None:
+        config = ModelConfig()
+    return load_model(config), load_tokenizer(config)

src/tuning/__init__.py ADDED Viewed

	@@ -0,0 +1,15 @@

+"""Tuning subpackage for Précis."""
+from src.tuning.lora import get_lora_config, apply_lora, merge_and_save
+from src.tuning.data import SummarizationDataset, prepare_dataset, create_dummy_data
+from src.tuning.trainer import PrecisTrainer
+__all__ = [
+    "get_lora_config",
+    "apply_lora",
+    "merge_and_save",
+    "SummarizationDataset",
+    "prepare_dataset",
+    "create_dummy_data",
+    "PrecisTrainer",
+]

src/tuning/data.py ADDED Viewed

	@@ -0,0 +1,70 @@

+"""Data preparation utilities for training."""
+import logging
+from typing import Dict, List, Optional, Any
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizer
+from src.config import DataConfig
+logger = logging.getLogger(__name__)
+class SummarizationDataset(Dataset):
+    """PyTorch Dataset for summarization training."""
+    def __init__(
+        self,
+        data: List[Dict[str, str]],
+        tokenizer: PreTrainedTokenizer,
+        config: Optional[DataConfig] = None,
+    ):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.config = config or DataConfig()
+    def __len__(self) -> int:
+        return len(self.data)
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        item = self.data[idx]
+        prompt = self.config.format_prompt(item[self.config.input_column])
+        full_text = prompt + item[self.config.target_column] + self.tokenizer.eos_token
+        encoding = self.tokenizer(
+            full_text,
+            truncation=True,
+            max_length=self.config.max_input_length + self.config.max_target_length,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "labels": encoding["input_ids"].squeeze(),
+        }
+def create_dummy_data(num_samples: int = 10) -> List[Dict[str, str]]:
+    """Generate dummy data for testing the training pipeline."""
+    samples = []
+    for i in range(num_samples):
+        samples.append({
+            "text": f"This is sample document {i}. It contains information about topic {i % 3}. "
+                    f"The document discusses various aspects and provides detailed analysis. "
+                    f"Key points include methodology, results, and conclusions for study {i}.",
+            "summary": f"Document {i} analyzes topic {i % 3}, covering methodology, results, and conclusions.",
+        })
+    logger.info(f"Created {num_samples} dummy samples")
+    return samples
+def prepare_dataset(
+    data: List[Dict[str, str]],
+    tokenizer: PreTrainedTokenizer,
+    config: Optional[DataConfig] = None,
+) -> SummarizationDataset:
+    """Prepare dataset for training."""
+    return SummarizationDataset(data, tokenizer, config)

src/tuning/lora.py ADDED Viewed

	@@ -0,0 +1,53 @@

+"""LoRA/PEFT configuration and utilities."""
+import logging
+from pathlib import Path
+from typing import Optional
+from peft import LoraConfig, get_peft_model, PeftModel, TaskType
+from transformers import PreTrainedModel
+from src.config import TrainingConfig
+logger = logging.getLogger(__name__)
+def get_lora_config(config: Optional[TrainingConfig] = None) -> LoraConfig:
+    """Create LoRA configuration for summarization task."""
+    if config is None:
+        config = TrainingConfig()
+    return LoraConfig(
+        r=config.lora_r,
+        lora_alpha=config.lora_alpha,
+        lora_dropout=config.lora_dropout,
+        target_modules=config.lora_target_modules,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+    )
+def apply_lora(model: PreTrainedModel, config: Optional[TrainingConfig] = None) -> PeftModel:
+    """Apply LoRA adapters to model."""
+    lora_config = get_lora_config(config)
+    logger.info(f"Applying LoRA with r={lora_config.r}, alpha={lora_config.lora_alpha}")
+    peft_model = get_peft_model(model, lora_config)
+    peft_model.print_trainable_parameters()
+    return peft_model
+def merge_and_save(model: PeftModel, output_path: str, tokenizer=None) -> None:
+    """Merge LoRA weights into base model and save."""
+    output_dir = Path(output_path)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    logger.info("Merging LoRA weights...")
+    merged_model = model.merge_and_unload()
+    logger.info(f"Saving merged model to {output_dir}")
+    merged_model.save_pretrained(output_dir)
+    if tokenizer:
+        tokenizer.save_pretrained(output_dir)

src/tuning/trainer.py ADDED Viewed

	@@ -0,0 +1,100 @@

+"""Training orchestration for Précis."""
+import logging
+from pathlib import Path
+from typing import Optional
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    DataCollatorForLanguageModeling,
+)
+from torch.utils.data import Dataset
+from src.config import TrainingConfig
+logger = logging.getLogger(__name__)
+class PrecisTrainer:
+    """Wrapper around HuggingFace Trainer for summarization fine-tuning."""
+    def __init__(
+        self,
+        model: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+        train_dataset: Dataset,
+        eval_dataset: Optional[Dataset] = None,
+        config: Optional[TrainingConfig] = None,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.train_dataset = train_dataset
+        self.eval_dataset = eval_dataset
+        self.config = config or TrainingConfig()
+        self.training_args = self._create_training_args()
+        self.trainer = self._create_trainer()
+    def _create_training_args(self) -> TrainingArguments:
+        """Create HuggingFace TrainingArguments from config."""
+        return TrainingArguments(
+            output_dir=self.config.output_dir,
+            num_train_epochs=self.config.num_epochs,
+            per_device_train_batch_size=self.config.batch_size,
+            gradient_accumulation_steps=self.config.gradient_accumulation_steps,
+            learning_rate=self.config.learning_rate,
+            warmup_ratio=self.config.warmup_ratio,
+            weight_decay=self.config.weight_decay,
+            max_grad_norm=self.config.max_grad_norm,
+            optim=self.config.optim,
+            logging_steps=self.config.logging_steps,
+            save_steps=self.config.save_steps,
+            eval_steps=self.config.eval_steps if self.eval_dataset else None,
+            evaluation_strategy="steps" if self.eval_dataset else "no",
+            save_total_limit=3,
+            load_best_model_at_end=bool(self.eval_dataset),
+            seed=self.config.seed,
+            fp16=True,
+            report_to="none",
+        )
+    def _create_trainer(self) -> Trainer:
+        """Create HuggingFace Trainer instance."""
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=self.tokenizer,
+            mlm=False,
+        )
+        return Trainer(
+            model=self.model,
+            args=self.training_args,
+            train_dataset=self.train_dataset,
+            eval_dataset=self.eval_dataset,
+            data_collator=data_collator,
+        )
+    def train(self) -> None:
+        """Execute training loop."""
+        logger.info("Starting training...")
+        self.trainer.train()
+        logger.info("Training complete.")
+    def evaluate(self) -> dict:
+        """Run evaluation and return metrics."""
+        if self.eval_dataset is None:
+            logger.warning("No eval dataset provided")
+            return {}
+        logger.info("Running evaluation...")
+        return self.trainer.evaluate()
+    def save(self, output_path: Optional[str] = None) -> None:
+        """Save model checkpoint."""
+        path = output_path or self.config.output_dir
+        Path(path).mkdir(parents=True, exist_ok=True)
+        self.trainer.save_model(path)
+        self.tokenizer.save_pretrained(path)
+        logger.info(f"Model saved to {path}")