LLM Engineer’s Handbook: Master the art of engineering large language models from concept to production

from zenml import pipeline, step

@step
def load_data() -> pd.DataFrame:
    """Load raw data from source."""
    return pd.read_csv("data.csv")

@step  
def preprocess_data(data: pd.DataFrame) -> pd.DataFrame:
    """Clean and preprocess data."""
    return data.dropna()

@step
def train_model(data: pd.DataFrame) -> Any:
    """Train ML model on processed data."""
    return model

@pipeline
def training_pipeline():
    """Complete training pipeline."""
    raw_data = load_data()
    clean_data = preprocess_data(raw_data)
    model = train_model(clean_data)
    return model

from qdrant_client.models import Filter, FieldCondition, MatchValue

# Search with metadata filtering
results = qdrant_client.search(
    collection_name="documents",
    query_vector=embedding,
    query_filter=Filter(
        must=[
            FieldCondition(
                key="author_id",
                match=MatchValue(value="author_123")
            )
        ]
    ),
    limit=10
)

# Create role with necessary permissions
poetry poe create-sagemaker-execution-role

# Update .env file with generated ARN
AWS_ARN_ROLE=arn:aws:iam::123456789012:role/SageMakerExecutionRole

from zenml import pipeline, step
from typing import Dict, Any

@step
def extract_data_step(source_config: Dict[str, Any]) -> List[Dict]:
    """Extract data from configured sources."""
    dispatcher = CrawlerDispatcher()
    return dispatcher.crawl_all_sources(source_config)

@step  
def transform_data_step(raw_data: List[Dict]) -> List[Dict]:
    """Transform and standardize extracted data."""
    cleaned_data = []
    for item in raw_data:
        cleaned_item = standardize_document(item)
        cleaned_data.append(cleaned_item)
    return cleaned_data

@step
def load_data_step(transformed_data: List[Dict]) -> bool:
    """Load data into warehouse."""
    warehouse = NoSQLDataWarehouse()
    return warehouse.bulk_insert(transformed_data)

@pipeline  
def data_collection_pipeline():
    """Complete data collection ETL pipeline."""
    raw_data = extract_data_step()
    clean_data = transform_data_step(raw_data)
    success = load_data_step(clean_data)
    return success

from abc import ABC, abstractmethod
from typing import List, Dict, Any
from selenium import webdriver

class BaseCrawler(ABC):
    """Abstract base class for all crawlers."""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.driver = None
        
    @abstractmethod
    def extract(self, source_url: str) -> List[Dict[str, Any]]:
        """Extract data from source URL."""
        pass
    
    def setup_driver(self) -> webdriver.Chrome:
        """Initialize Selenium WebDriver with optimal settings."""
        options = webdriver.ChromeOptions()
        options.add_argument('--headless')
        options.add_argument('--no-sandbox')
        options.add_argument('--disable-dev-shm-usage')
        return webdriver.Chrome(options=options)
    
    def cleanup(self):
        """Clean up resources."""
        if self.driver:
            self.driver.quit()

class GitHubCrawler(BaseCrawler):
    """Crawler for GitHub repositories and documentation."""
    
    def extract(self, repo_url: str) -> List[Dict[str, Any]]:
        """Extract README and documentation files."""
        documents = []
        
        # Clone repository or use GitHub API
        repo_data = self._fetch_repo_content(repo_url)
        
        for file_path, content in repo_data.items():
            if self._is_documentation_file(file_path):
                doc = {
                    'content': content,
                    'source_url': f"{repo_url}/blob/main/{file_path}",
                    'platform': 'github',
                    'file_type': self._get_file_type(file_path),
                    'created_at': self._get_file_timestamp(file_path),
                    'metadata': {
                        'repository': repo_url,
                        'file_path': file_path,
                        'language': self._detect_language(content)
                    }
                }
                documents.append(doc)
                
        return documents
    
    def _is_documentation_file(self, file_path: str) -> bool:
        """Check if file contains documentation."""
        doc_patterns = ['.md', '.txt', '.rst', 'README', 'CHANGELOG']
        return any(pattern in file_path for pattern in doc_patterns)

class MediumCrawler(BaseCrawler):
    """Crawler for Medium articles and publications."""
    
    def extract(self, profile_url: str) -> List[Dict[str, Any]]:
        """Extract articles from Medium profile."""
        self.driver = self.setup_driver()
        articles = []
        
        try:
            # Navigate to profile page
            self.driver.get(profile_url)
            
            # Scroll to load all articles
            self._scroll_to_load_content()
            
            # Extract article links
            article_links = self._get_article_links()
            
            # Extract content from each article
            for link in article_links:
                article_data = self._extract_article_content(link)
                if article_data:
                    articles.append(article_data)
                    
        finally:
            self.cleanup()
            
        return articles
    
    def _extract_article_content(self, article_url: str) -> Dict[str, Any]:
        """Extract content from individual Medium article."""
        self.driver.get(article_url)
        
        # Wait for content to load
        WebDriverWait(self.driver, 10).until(
            EC.presence_of_element_located((By.TAG_NAME, "article"))
        )
        
        # Extract article data
        title = self.driver.find_element(By.TAG_NAME, "h1").text
        content = self._extract_article_text()
        published_date = self._extract_publish_date()
        
        return {
            'title': title,
            'content': content,
            'source_url': article_url,
            'platform': 'medium',
            'published_at': published_date,
            'metadata': {
                'word_count': len(content.split()),
                'reading_time': self._estimate_reading_time(content)
            }
        }

class CrawlerDispatcher:
    """Dispatcher to route crawling requests to appropriate crawlers."""
    
    def __init__(self):
        self.crawlers = {
            'github': GitHubCrawler,
            'medium': MediumCrawler,
            'linkedin': LinkedInCrawler,
            'custom': CustomArticleCrawler
        }
    
    def get_crawler(self, platform: str, config: Dict[str, Any]) -> BaseCrawler:
        """Factory method to instantiate the right crawler."""
        if platform not in self.crawlers:
            raise ValueError(f"Unsupported platform: {platform}")
            
        crawler_class = self.crawlers[platform]
        return crawler_class(config)
    
    def crawl_all_sources(self, sources_config: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Crawl all configured data sources."""
        all_documents = []
        
        for platform, config in sources_config.items():
            try:
                crawler = self.get_crawler(platform, config)
                documents = crawler.extract(config['source_url'])
                all_documents.extend(documents)
                
            except Exception as e:
                logger.error(f"Failed to crawl {platform}: {str(e)}")
                continue
                
        return all_documents

from pydantic import BaseModel, Field
from datetime import datetime
from typing import Optional, Dict, Any
from uuid import UUID, uuid4

class UserDocument(BaseModel):
    """User profile document schema."""
    
    id: UUID = Field(default_factory=uuid4)
    full_name: str
    email: Optional[str] = None
    bio: Optional[str] = None
    linkedin_url: Optional[str] = None
    medium_url: Optional[str] = None
    github_url: Optional[str] = None
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)
    
    class Config:
        json_encoders = {
            datetime: lambda v: v.isoformat(),
            UUID: lambda v: str(v)
        }

class BaseDocument(BaseModel):
    """Base schema for all content documents."""
    
    id: UUID = Field(default_factory=uuid4)
    content: str
    title: Optional[str] = None
    source_url: str
    platform: str  # github, medium, linkedin, custom
    author_id: UUID
    created_at: datetime = Field(default_factory=datetime.utcnow)
    metadata: Dict[str, Any] = Field(default_factory=dict)

class ArticleDocument(BaseDocument):
    """Schema for articles and blog posts."""
    
    published_at: Optional[datetime] = None
    word_count: Optional[int] = None
    reading_time: Optional[int] = None  # minutes
    tags: List[str] = Field(default_factory=list)

class PostDocument(BaseDocument):
    """Schema for social media posts."""
    
    engagement_metrics: Optional[Dict[str, int]] = None
    hashtags: List[str] = Field(default_factory=list)
    mentions: List[str] = Field(default_factory=list)

class RepositoryDocument(BaseDocument):
    """Schema for code repositories."""
    
    language: Optional[str] = None
    file_path: str
    file_type: str  # .md, .py, .txt, etc.
    repository_url: str
    branch: str = "main"

from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import IndexModel, ASCENDING, TEXT

class NoSQLDataWarehouse:
    """MongoDB-based data warehouse with ODM capabilities."""
    
    def __init__(self, connection_string: str, database_name: str):
        self.client = AsyncIOMotorClient(connection_string)
        self.db = self.client[database_name]
        self._ensure_indexes()
    
    async def _ensure_indexes(self):
        """Create necessary indexes for optimal query performance."""
        # User collection indexes
        await self.db.users.create_indexes([
            IndexModel([("email", ASCENDING)], unique=True),
            IndexModel([("full_name", TEXT)])
        ])
        
        # Content collection indexes  
        for collection_name in ['articles', 'posts', 'repositories']:
            collection = self.db[collection_name]
            await collection.create_indexes([
                IndexModel([("author_id", ASCENDING)]),
                IndexModel([("platform", ASCENDING)]),
                IndexModel([("created_at", ASCENDING)]),
                IndexModel([("content", TEXT), ("title", TEXT)])
            ])
    
    async def insert_document(self, document: BaseDocument, collection_name: str) -> str:
        """Insert a single document."""
        collection = self.db[collection_name]
        doc_dict = document.dict()
        result = await collection.insert_one(doc_dict)
        return str(result.inserted_id)
    
    async def bulk_insert_documents(self, documents: List[BaseDocument], 
                                  collection_name: str) -> List[str]:
        """Bulk insert multiple documents."""
        if not documents:
            return []
            
        collection = self.db[collection_name]
        doc_dicts = [doc.dict() for doc in documents]
        result = await collection.insert_many(doc_dicts)
        return [str(oid) for oid in result.inserted_ids]
    
    async def find_documents_by_author(self, author_id: str, 
                                     collection_name: str) -> List[Dict[str, Any]]:
        """Retrieve all documents for a specific author."""
        collection = self.db[collection_name]
        cursor = collection.find({"author_id": author_id})
        return await cursor.to_list(length=None)

import re
from typing import List, Dict, Any
from bs4 import BeautifulSoup

class DataCleaner:
    """Comprehensive data cleaning utilities."""
    
    @staticmethod
    def clean_html_content(content: str) -> str:
        """Remove HTML tags and clean text."""
        # Parse HTML and extract text
        soup = BeautifulSoup(content, 'html.parser')
        
        # Remove script and style elements
        for script in soup(["script", "style"]):
            script.decompose()
        
        # Get text and clean whitespace
        text = soup.get_text()
        text = re.sub(r'\s+', ' ', text)  # Normalize whitespace
        text = text.strip()
        
        return text
    
    @staticmethod
    def remove_personal_info(content: str) -> str:
        """Remove or mask personal information."""
        # Remove email addresses
        content = re.sub(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b', 
                        '[EMAIL]', content)
        
        # Remove phone numbers (basic pattern)
        content = re.sub(r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b', '[PHONE]', content)
        
        # Remove URLs
        content = re.sub(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', 
                        '[URL]', content)
        
        return content
    
    @staticmethod
    def standardize_encoding(content: str) -> str:
        """Standardize text encoding and fix common issues."""
        # Fix common encoding issues
        replacements = {
            ''': "'",
            ''': "'", 
            '"': '"',
            '"': '"',
            '–': '-',
            '—': '-',
            '…': '...'
        }
        
        for old, new in replacements.items():
            content = content.replace(old, new)
        
        return content
    
    @classmethod
    def clean_document(cls, document: Dict[str, Any]) -> Dict[str, Any]:
        """Apply all cleaning steps to a document."""
        content = document.get('content', '')
        
        # Apply cleaning steps
        content = cls.clean_html_content(content)
        content = cls.remove_personal_info(content)
        content = cls.standardize_encoding(content)
        
        # Update document
        document['content'] = content
        document['metadata']['cleaned'] = True
        document['metadata']['original_length'] = len(document.get('content', ''))
        document['metadata']['cleaned_length'] = len(content)
        
        return document

import asyncio
from tenacity import retry, stop_after_attempt, wait_exponential

class RobustCrawler:
    """Crawler with comprehensive error handling."""
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10)
    )
    async def crawl_with_retry(self, url: str) -> Dict[str, Any]:
        """Crawl URL with automatic retry logic."""
        try:
            return await self._crawl_url(url)
        except TimeoutException:
            logger.warning(f"Timeout crawling {url}, retrying...")
            raise
        except ConnectionError:
            logger.warning(f"Connection error for {url}, retrying...")
            raise
        except Exception as e:
            logger.error(f"Unexpected error crawling {url}: {str(e)}")
            raise
    
    async def crawl_with_fallback(self, primary_url: str, 
                                 fallback_urls: List[str]) -> Optional[Dict[str, Any]]:
        """Try primary URL, fall back to alternatives if it fails."""
        urls_to_try = [primary_url] + fallback_urls
        
        for url in urls_to_try:
            try:
                return await self.crawl_with_retry(url)
            except Exception as e:
                logger.warning(f"Failed to crawl {url}: {str(e)}")
                continue
        
        logger.error(f"Failed to crawl any URL for source")
        return None

from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List

class EmbeddingGenerator:
    """Generate and manage document embeddings."""
    
    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.embedding_dim = self.model.get_sentence_embedding_dimension()
    
    def generate_embeddings(self, texts: List[str]) -> np.ndarray:
        """Generate embeddings for a list of texts."""
        embeddings = self.model.encode(
            texts,
            batch_size=32,
            show_progress_bar=True,
            normalize_embeddings=True  # L2 normalization for cosine similarity
        )
        return embeddings
    
    def calculate_similarity(self, embedding1: np.ndarray, 
                           embedding2: np.ndarray) -> float:
        """Calculate cosine similarity between two embeddings."""
        return np.dot(embedding1, embedding2) / (
            np.linalg.norm(embedding1) * np.linalg.norm(embedding2)
        )
    
    def find_most_similar(self, query_embedding: np.ndarray, 
                         candidate_embeddings: np.ndarray, 
                         top_k: int = 5) -> List[int]:
        """Find indices of most similar embeddings."""
        similarities = np.dot(candidate_embeddings, query_embedding)
        top_indices = np.argsort(similarities)[::-1][:top_k]
        return top_indices.tolist()

from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct, Filter, 
    FieldCondition, MatchValue
)
from typing import List, Dict, Any, Optional
import uuid

class VectorDBManager:
    """Manage vector database operations with Qdrant."""
    
    def __init__(self, host: str = "localhost", port: int = 6333):
        self.client = QdrantClient(host=host, port=port)
        
    def create_collection(self, collection_name: str, 
                         vector_size: int = 384,
                         distance: Distance = Distance.COSINE):
        """Create a new collection with specified parameters."""
        self.client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(
                size=vector_size,
                distance=distance
            )
        )
    
    def upsert_documents(self, collection_name: str, 
                        documents: List[Dict[str, Any]],
                        embeddings: np.ndarray) -> bool:
        """Insert or update documents with their embeddings."""
        points = []
        
        for i, doc in enumerate(documents):
            point = PointStruct(
                id=str(uuid.uuid4()),
                vector=embeddings[i].tolist(),
                payload={
                    "content": doc["content"],
                    "title": doc.get("title", ""),
                    "source_url": doc["source_url"],
                    "platform": doc["platform"],
                    "author_id": str(doc["author_id"]),
                    "metadata": doc.get("metadata", {})
                }
            )
            points.append(point)
        
        operation_result = self.client.upsert(
            collection_name=collection_name,
            points=points
        )
        
        return operation_result.status == "completed"
    
    def search_similar_documents(self, collection_name: str,
                               query_vector: List[float],
                               top_k: int = 10,
                               author_filter: Optional[str] = None) -> List[Dict]:
        """Search for similar documents with optional filtering."""
        search_filter = None
        
        if author_filter:
            search_filter = Filter(
                must=[
                    FieldCondition(
                        key="author_id",
                        match=MatchValue(value=author_filter)
                    )
                ]
            )
        
        search_results = self.client.search(
            collection_name=collection_name,
            query_vector=query_vector,
            query_filter=search_filter,
            limit=top_k,
            with_payload=True
        )
        
        return [
            {
                "id": result.id,
                "score": result.score,
                "content": result.payload["content"],
                "title": result.payload.get("title", ""),
                "source_url": result.payload["source_url"],
                "platform": result.payload["platform"],
                "metadata": result.payload.get("metadata", {})
            }
            for result in search_results
        ]

from zenml import pipeline, step
from typing import List, Dict, Any, Tuple
import pandas as pd

@step
def query_data_warehouse_step(
    query_params: Dict[str, Any]
) -> pd.DataFrame:
    """Query the MongoDB data warehouse for documents to process."""
    warehouse = NoSQLDataWarehouse()
    
    # Query for new or updated documents
    documents = warehouse.query_documents_since_last_run(
        last_run_timestamp=query_params.get("last_run_timestamp"),
        author_ids=query_params.get("author_ids", []),
        platforms=query_params.get("platforms", [])
    )
    
    return pd.DataFrame(documents)

@step
def clean_documents_step(raw_documents: pd.DataFrame) -> pd.DataFrame:
    """Clean and preprocess document content."""
    cleaner = DocumentCleaner()
    
    cleaned_docs = []
    for _, doc in raw_documents.iterrows():
        cleaned_doc = cleaner.clean_document(doc.to_dict())
        cleaned_docs.append(cleaned_doc)
    
    return pd.DataFrame(cleaned_docs)

@step  
def chunk_documents_step(
    clean_documents: pd.DataFrame,
    chunk_size: int = 500,
    chunk_overlap: int = 50
) -> pd.DataFrame:
    """Split documents into smaller chunks for embedding."""
    chunker = DocumentChunker(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap
    )
    
    all_chunks = []
    for _, doc in clean_documents.iterrows():
        chunks = chunker.chunk_document(doc.to_dict())
        all_chunks.extend(chunks)
    
    return pd.DataFrame(all_chunks)

@step
def embed_chunks_step(
    chunks: pd.DataFrame,
    model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
) -> Tuple[pd.DataFrame, np.ndarray]:
    """Generate embeddings for document chunks."""
    embedder = EmbeddingGenerator(model_name)
    
    # Extract content for embedding
    texts = chunks['content'].tolist()
    
    # Generate embeddings
    embeddings = embedder.generate_embeddings(texts)
    
    # Add embedding metadata
    chunks['embedding_model'] = model_name
    chunks['embedding_dim'] = embeddings.shape[1]
    
    return chunks, embeddings

@step
def load_to_vector_db_step(
    chunks: pd.DataFrame, 
    embeddings: np.ndarray,
    collection_name: str = "document_chunks"
) -> bool:
    """Load chunks and embeddings to vector database."""
    vector_db = VectorDBManager()
    
    # Ensure collection exists
    try:
        vector_db.create_collection(
            collection_name=collection_name,
            vector_size=embeddings.shape[1]
        )
    except Exception:
        # Collection might already exist
        pass
    
    # Convert DataFrame to list of dicts
    chunk_dicts = chunks.to_dict('records')
    
    # Upsert to vector database
    success = vector_db.upsert_documents(
        collection_name=collection_name,
        documents=chunk_dicts,
        embeddings=embeddings
    )
    
    return success

@pipeline
def rag_feature_pipeline():
    """Complete RAG feature processing pipeline."""
    # Query new/updated documents
    raw_docs = query_data_warehouse_step()
    
    # Clean and preprocess
    clean_docs = clean_documents_step(raw_docs)
    
    # Chunk documents
    chunks = chunk_documents_step(clean_docs)
    
    # Generate embeddings
    chunks_with_metadata, embeddings = embed_chunks_step(chunks)
    
    # Load to vector database
    success = load_to_vector_db_step(chunks_with_metadata, embeddings)
    
    return success

class DocumentChunker:
    """Advanced document chunking with multiple strategies."""
    
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
    
    def chunk_by_tokens(self, text: str) -> List[str]:
        """Chunk document by token count."""
        import tiktoken
        
        encoder = tiktoken.get_encoding("cl100k_base")
        tokens = encoder.encode(text)
        
        chunks = []
        for i in range(0, len(tokens), self.chunk_size - self.chunk_overlap):
            chunk_tokens = tokens[i:i + self.chunk_size]
            chunk_text = encoder.decode(chunk_tokens)
            chunks.append(chunk_text)
        
        return chunks
    
    def chunk_by_sentences(self, text: str) -> List[str]:
        """Chunk document by sentence boundaries."""
        import nltk
        nltk.download('punkt', quiet=True)
        
        sentences = nltk.sent_tokenize(text)
        chunks = []
        current_chunk = []
        current_length = 0
        
        for sentence in sentences:
            sentence_length = len(sentence.split())
            
            if current_length + sentence_length > self.chunk_size and current_chunk:
                # Start new chunk
                chunks.append(' '.join(current_chunk))
                
                # Handle overlap
                overlap_sentences = current_chunk[-self.chunk_overlap//20:] if self.chunk_overlap > 0 else []
                current_chunk = overlap_sentences + [sentence]
                current_length = sum(len(s.split()) for s in current_chunk)
            else:
                current_chunk.append(sentence)
                current_length += sentence_length
        
        # Add final chunk
        if current_chunk:
            chunks.append(' '.join(current_chunk))
        
        return chunks
    
    def chunk_by_semantic_similarity(self, text: str, threshold: float = 0.7) -> List[str]:
        """Chunk document by semantic coherence."""
        sentences = nltk.sent_tokenize(text)
        embedder = EmbeddingGenerator()
        
        # Generate embeddings for all sentences
        sentence_embeddings = embedder.generate_embeddings(sentences)
        
        chunks = []
        current_chunk = [sentences[0]]
        current_embeddings = [sentence_embeddings[0]]
        
        for i in range(1, len(sentences)):
            # Calculate similarity with current chunk
            chunk_centroid = np.mean(current_embeddings, axis=0)
            similarity = embedder.calculate_similarity(
                sentence_embeddings[i], 
                chunk_centroid
            )
            
            if similarity >= threshold or len(current_chunk) < 3:
                # Add to current chunk
                current_chunk.append(sentences[i])
                current_embeddings.append(sentence_embeddings[i])
            else:
                # Start new chunk
                chunks.append(' '.join(current_chunk))
                current_chunk = [sentences[i]]
                current_embeddings = [sentence_embeddings[i]]
        
        # Add final chunk
        if current_chunk:
            chunks.append(' '.join(current_chunk))
        
        return chunks
    
    def chunk_document(self, document: Dict[str, Any], 
                      strategy: str = "tokens") -> List[Dict[str, Any]]:
        """Chunk a document using specified strategy."""
        content = document['content']
        
        # Choose chunking strategy
        if strategy == "tokens":
            chunks = self.chunk_by_tokens(content)
        elif strategy == "sentences":
            chunks = self.chunk_by_sentences(content)
        elif strategy == "semantic":
            chunks = self.chunk_by_semantic_similarity(content)
        else:
            raise ValueError(f"Unknown chunking strategy: {strategy}")
        
        # Create chunk documents
        chunk_documents = []
        for i, chunk in enumerate(chunks):
            chunk_doc = {
                **document,  # Copy all original fields
                'content': chunk,
                'chunk_index': i,
                'total_chunks': len(chunks),
                'chunk_strategy': strategy,
                'parent_document_id': document['id']
            }
            chunk_documents.append(chunk_doc)
        
        return chunk_documents

def sliding_window_chunk(text: str, window_size: int = 500, 
                        stride: int = 250) -> List[Dict[str, Any]]:
    """Create overlapping chunks with sliding window."""
    words = text.split()
    chunks = []
    
    for i in range(0, len(words), stride):
        chunk_words = words[i:i + window_size]
        if len(chunk_words) < window_size // 2:  # Skip very small chunks
            break
            
        chunk = {
            'content': ' '.join(chunk_words),
            'start_index': i,
            'end_index': min(i + window_size, len(words)),
            'overlap_prev': i > 0,
            'overlap_next': i + window_size < len(words)
        }
        chunks.append(chunk)
    
    return chunks

def enhance_with_metadata(document: Dict[str, Any]) -> Dict[str, Any]:
    """Add rich metadata for better filtering and retrieval."""
    content = document['content']
    
    # Extract structural metadata
    metadata = {
        'word_count': len(content.split()),
        'char_count': len(content),
        'paragraph_count': len([p for p in content.split('\n\n') if p.strip()]),
        'avg_sentence_length': calculate_avg_sentence_length(content),
        'readability_score': calculate_readability(content),
        'language': detect_language(content),
        'topics': extract_topics(content),
        'entities': extract_named_entities(content),
        'sentiment': analyze_sentiment(content)
    }
    
    document['metadata'].update(metadata)
    return document

class HybridRetriever:
    """Combine dense vector search with sparse keyword search."""
    
    def __init__(self, vector_db: VectorDBManager, elastic_client):
        self.vector_db = vector_db
        self.elastic = elastic_client
        
    def hybrid_search(self, query: str, collection_name: str,
                     top_k: int = 10, alpha: float = 0.7) -> List[Dict]:
        """Combine vector and keyword search results."""
        # Dense vector search
        query_embedding = self.embed_query(query)
        vector_results = self.vector_db.search_similar_documents(
            collection_name=collection_name,
            query_vector=query_embedding,
            top_k=top_k * 2  # Get more candidates
        )
        
        # Sparse keyword search
        keyword_results = self.elastic.search(
            index=collection_name,
            body={
                "query": {
                    "multi_match": {
                        "query": query,
                        "fields": ["content", "title"],
                        "type": "best_fields"
                    }
                },
                "size": top_k * 2
            }
        )
        
        # Combine and rerank results
        combined_results = self._combine_results(
            vector_results, keyword_results, alpha
        )
        
        return combined_results[:top_k]
    
    def _combine_results(self, vector_results: List[Dict], 
                        keyword_results: Dict, alpha: float) -> List[Dict]:
        """Combine and score results from both search methods."""
        # Normalize scores and combine
        result_map = {}
        
        # Process vector results
        max_vector_score = max(r['score'] for r in vector_results) if vector_results else 1
        for result in vector_results:
            doc_id = result['id']
            normalized_score = result['score'] / max_vector_score
            result_map[doc_id] = {
                **result,
                'vector_score': normalized_score,
                'keyword_score': 0.0,
                'combined_score': alpha * normalized_score
            }
        
        # Process keyword results
        hits = keyword_results['hits']['hits']
        max_keyword_score = max(hit['_score'] for hit in hits) if hits else 1
        
        for hit in hits:
            doc_id = hit['_id']
            normalized_score = hit['_score'] / max_keyword_score
            
            if doc_id in result_map:
                # Update existing result
                result_map[doc_id]['keyword_score'] = normalized_score
                result_map[doc_id]['combined_score'] = (
                    alpha * result_map[doc_id]['vector_score'] + 
                    (1 - alpha) * normalized_score
                )
            else:
                # Add new result
                result_map[doc_id] = {
                    'id': doc_id,
                    'content': hit['_source']['content'],
                    'title': hit['_source'].get('title', ''),
                    'vector_score': 0.0,
                    'keyword_score': normalized_score,
                    'combined_score': (1 - alpha) * normalized_score
                }
        
        # Sort by combined score
        sorted_results = sorted(
            result_map.values(),
            key=lambda x: x['combined_score'],
            reverse=True
        )
        
        return sorted_results

from pymongo.change_stream import ChangeStream
import asyncio

class CDCManager:
    """Manage change data capture from MongoDB to trigger pipeline updates."""
    
    def __init__(self, mongo_client, vector_db: VectorDBManager):
        self.mongo_client = mongo_client
        self.vector_db = vector_db
        self.running = False
    
    async def watch_collection_changes(self, collection_name: str):
        """Watch for changes in MongoDB collection and update vector DB."""
        collection = self.mongo_client.get_collection(collection_name)
        
        # Define pipeline to watch for insert, update, delete operations
        pipeline = [
            {
                '$match': {
                    'operationType': {'$in': ['insert', 'update', 'delete']}
                }
            }
        ]
        
        self.running = True
        
        try:
            with collection.watch(pipeline) as change_stream:
                for change in change_stream:
                    if not self.running:
                        break
                        
                    await self._process_change(change, collection_name)
                    
        except Exception as e:
            logger.error(f"CDC error for {collection_name}: {str(e)}")
        
    async def _process_change(self, change: Dict[str, Any], collection_name: str):
        """Process individual change events."""
        operation_type = change['operationType']
        
        if operation_type == 'insert':
            await self._handle_insert(change, collection_name)
        elif operation_type == 'update':
            await self._handle_update(change, collection_name)
        elif operation_type == 'delete':
            await self._handle_delete(change, collection_name)
    
    async def _handle_insert(self, change: Dict[str, Any], collection_name: str):
        """Handle document insertion."""
        document = change['fullDocument']
        
        # Process new document through RAG pipeline
        chunks = self._chunk_document(document)
        embeddings = self._generate_embeddings(chunks)
        
        # Update vector database
        await self.vector_db.upsert_documents(
            collection_name=f"{collection_name}_chunks",
            documents=chunks,
            embeddings=embeddings
        )
        
        logger.info(f"Processed new document: {document.get('id', 'unknown')}")
    
    async def _handle_update(self, change: Dict[str, Any], collection_name: str):
        """Handle document updates."""
        document_id = change['documentKey']['_id']
        updated_fields = change.get('updateDescription', {}).get('updatedFields', {})
        
        # If content was updated, reprocess the document
        if 'content' in updated_fields:
            # Remove old chunks
            await self._remove_document_chunks(document_id, collection_name)
            
            # Process updated document
            full_document = await self._fetch_full_document(document_id, collection_name)
            await self._handle_insert({'fullDocument': full_document}, collection_name)
        
        logger.info(f"Processed document update: {document_id}")
    
    async def _handle_delete(self, change: Dict[str, Any], collection_name: str):
        """Handle document deletion."""
        document_id = change['documentKey']['_id']
        
        # Remove associated chunks from vector database
        await self._remove_document_chunks(document_id, collection_name)
        
        logger.info(f"Processed document deletion: {document_id}")
    
    def stop_watching(self):
        """Stop the CDC process."""
        self.running = False

from typing import Dict, Any, List, Optional
from pydantic import BaseModel

class InstructionSample(BaseModel):
    """Schema for instruction dataset samples."""
    
    instruction: str  # Task description or question
    input: Optional[str] = None  # Additional context or data
    output: str  # Expected response or answer
    system: Optional[str] = None  # System prompt for behavior steering
    
    class Config:
        schema_extra = {
            "example": {
                "system": "You are a helpful AI assistant that provides clear explanations.",
                "instruction": "Explain the concept of supervised fine-tuning for large language models.",
                "input": "Focus on the key benefits and use cases.",
                "output": "Supervised fine-tuning (SFT) is a process where a pre-trained language model..."
            }
        }

import re
from typing import List, Dict, Any
from langdetect import detect

class InstructionDatasetFilter:
    """Comprehensive filtering for instruction datasets."""
    
    def __init__(self):
        self.min_instruction_length = 10
        self.max_instruction_length = 2000
        self.min_output_length = 5
        self.max_output_length = 4000
        
        # Content quality patterns
        self.spam_patterns = [
            r'click here',
            r'subscribe now',
            r'buy now',
            r'\$\d+\.\d+',  # Prices
            r'www\.[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}'  # URLs
        ]
        
        self.low_quality_patterns = [
            r'^(yes|no|ok|sure)\.?$',  # Too short responses
            r'(.)\1{4,}',  # Repeated characters
            r'^[A-Z\s!]+$',  # All caps
        ]
    
    def filter_dataset(self, samples: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Apply all filters to dataset."""
        filtered_samples = []
        
        for sample in samples:
            if self._is_valid_sample(sample):
                filtered_samples.append(sample)
        
        return filtered_samples
    
    def _is_valid_sample(self, sample: Dict[str, Any]) -> bool:
        """Check if sample meets quality criteria."""
        instruction = sample.get('instruction', '')
        output = sample.get('output', '')
        
        # Length checks
        if not self._check_length(instruction, self.min_instruction_length, 
                                 self.max_instruction_length):
            return False
            
        if not self._check_length(output, self.min_output_length, 
                                 self.max_output_length):
            return False
        
        # Language detection
        try:
            if detect(instruction) != 'en' or detect(output) != 'en':
                return False
        except:
            return False  # Skip if language detection fails
        
        # Content quality checks
        if self._contains_spam(instruction) or self._contains_spam(output):
            return False
            
        if self._is_low_quality(output):
            return False
        
        # Instruction-output consistency
        if not self._check_instruction_output_consistency(instruction, output):
            return False
        
        return True
    
    def _check_length(self, text: str, min_len: int, max_len: int) -> bool:
        """Check if text length is within bounds."""
        return min_len <= len(text.split()) <= max_len
    
    def _contains_spam(self, text: str) -> bool:
        """Check for spam patterns."""
        text_lower = text.lower()
        return any(re.search(pattern, text_lower) for pattern in self.spam_patterns)
    
    def _is_low_quality(self, text: str) -> bool:
        """Check for low-quality content patterns."""
        return any(re.search(pattern, text, re.IGNORECASE) 
                  for pattern in self.low_quality_patterns)
    
    def _check_instruction_output_consistency(self, instruction: str, output: str) -> bool:
        """Basic consistency check between instruction and output."""
        # Check if output is not just repeating the instruction
        instruction_words = set(instruction.lower().split())
        output_words = set(output.lower().split())
        
        # If output is >80% same words as instruction, likely low quality
        if len(instruction_words & output_words) / len(output_words) > 0.8:
            return False
        
        return True

import hashlib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

class DataDeduplicator:
    """Remove duplicate and near-duplicate samples."""
    
    def __init__(self, similarity_threshold: float = 0.85):
        self.similarity_threshold = similarity_threshold
        self.vectorizer = TfidfVectorizer(
            ngram_range=(1, 3),
            max_features=10000,
            stop_words='english'
        )
    
    def deduplicate_dataset(self, samples: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove duplicates using multiple methods."""
        # 1. Exact deduplication
        samples = self._exact_deduplication(samples)
        
        # 2. Near-duplicate detection
        samples = self._fuzzy_deduplication(samples)
        
        return samples
    
    def _exact_deduplication(self, samples: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove exact duplicates using content hashing."""
        seen_hashes = set()
        unique_samples = []
        
        for sample in samples:
            # Create content hash from instruction + output
            content = sample['instruction'] + sample['output']
            content_hash = hashlib.md5(content.encode()).hexdigest()
            
            if content_hash not in seen_hashes:
                seen_hashes.add(content_hash)
                unique_samples.append(sample)
        
        return unique_samples
    
    def _fuzzy_deduplication(self, samples: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove near-duplicates using TF-IDF similarity."""
        if len(samples) < 2:
            return samples
        
        # Combine instruction and output for similarity calculation
        texts = [sample['instruction'] + ' ' + sample['output'] for sample in samples]
        
        # Vectorize texts
        tfidf_matrix = self.vectorizer.fit_transform(texts)
        
        # Calculate pairwise similarities
        similarity_matrix = cosine_similarity(tfidf_matrix)
        
        # Find duplicates
        to_remove = set()
        for i in range(len(similarity_matrix)):
            if i in to_remove:
                continue
                
            for j in range(i + 1, len(similarity_matrix)):
                if j in to_remove:
                    continue
                    
                if similarity_matrix[i][j] >= self.similarity_threshold:
                    # Keep the shorter sample (often higher quality)
                    if len(texts[i]) <= len(texts[j]):
                        to_remove.add(j)
                    else:
                        to_remove.add(i)
                        break
        
        # Return samples not marked for removal
        return [sample for i, sample in enumerate(samples) if i not in to_remove]

from openai import OpenAI
from typing import List, Dict, Any
import asyncio
import json

class InstructionDataGenerator:
    """Generate instruction-answer pairs from raw content."""
    
    def __init__(self, openai_client: OpenAI):
        self.client = openai_client
        self.generation_prompt = """
        Based on the following content, create a high-quality instruction-answer pair that would help someone learn from this information.

        Content: {content}

        Requirements:
        1. Create a clear, specific instruction that asks for information contained in the content
        2. Provide a comprehensive answer based solely on the content
        3. Make the instruction educational and valuable
        4. Ensure the answer is accurate and well-structured

        Format your response as JSON:
        {
            "instruction": "Your instruction here",
            "output": "Your detailed answer here"
        }
        """
    
    async def generate_instruction_pairs(self, content_documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Generate instruction-answer pairs from content documents."""
        instruction_pairs = []
        
        for doc in content_documents:
            try:
                pairs = await self._generate_pairs_for_document(doc)
                instruction_pairs.extend(pairs)
            except Exception as e:
                print(f"Error generating pairs for document {doc.get('id', 'unknown')}: {e}")
                continue
        
        return instruction_pairs
    
    async def _generate_pairs_for_document(self, document: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Generate multiple instruction-answer pairs from a single document."""
        content = document['content']
        
        # Split long content into chunks
        chunks = self._split_content_for_generation(content)
        
        pairs = []
        for chunk in chunks:
            try:
                response = await self.client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=[
                        {
                            "role": "system",
                            "content": "You are an expert at creating educational instruction-answer pairs from content."
                        },
                        {
                            "role": "user",
                            "content": self.generation_prompt.format(content=chunk)
                        }
                    ],
                    temperature=0.7,
                    max_tokens=1000
                )
                
                # Parse JSON response
                generated_content = response.choices[0].message.content
                pair_data = json.loads(generated_content)
                
                # Add metadata from original document
                pair = {
                    'instruction': pair_data['instruction'],
                    'output': pair_data['output'],
                    'source_document_id': document.get('id'),
                    'source_url': document.get('source_url'),
                    'platform': document.get('platform'),
                    'author_id': document.get('author_id'),
                    'generated_at': datetime.utcnow().isoformat()
                }
                
                pairs.append(pair)
                
            except Exception as e:
                print(f"Error generating pair for chunk: {e}")
                continue
        
        return pairs
    
    def _split_content_for_generation(self, content: str, max_chunk_size: int = 2000) -> List[str]:
        """Split content into manageable chunks for generation."""
        words = content.split()
        chunks = []
        
        current_chunk = []
        current_size = 0
        
        for word in words:
            if current_size + len(word) > max_chunk_size and current_chunk:
                chunks.append(' '.join(current_chunk))
                current_chunk = [word]
                current_size = len(word)
            else:
                current_chunk.append(word)
                current_size += len(word) + 1  # +1 for space
        
        if current_chunk:
            chunks.append(' '.join(current_chunk))
        
        return chunks

class ChatTemplateManager:
    """Manage different chat templates for fine-tuning."""
    
    TEMPLATES = {
        'alpaca': {
            'template': "### Instruction:\n{instruction}\n\n### Response:\n{response}",
            'eos_token': "</s>"
        },
        'chatml': {
            'template': "<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n{response}<|im_end|>",
            'eos_token': "<|im_end|>"
        },
        'llama3': {
            'template': "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{instruction}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n{response}<|eot_id|>",
            'eos_token': "<|eot_id|>"
        }
    }
    
    def format_sample(self, sample: Dict[str, Any], template_name: str) -> str:
        """Format instruction sample with specified template."""
        template_config = self.TEMPLATES[template_name]
        
        formatted = template_config['template'].format(
            instruction=sample['instruction'],
            response=sample['output']
        )
        
        # Add EOS token
        formatted += template_config['eos_token']
        
        return formatted
    
    def prepare_dataset_for_training(self, dataset: List[Dict[str, Any]], 
                                   template_name: str) -> List[str]:
        """Prepare entire dataset with consistent formatting."""
        return [self.format_sample(sample, template_name) for sample in dataset]

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM
import torch

class LoRATrainer:
    """LoRA fine-tuning implementation."""
    
    def __init__(self, model_name: str, lora_config: Dict[str, Any]):
        self.model_name = model_name
        self.lora_config = lora_config
        
    def setup_lora_model(self) -> tuple:
        """Setup model with LoRA adapters."""
        # Load base model
        model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        
        # Configure LoRA
        lora_config = LoraConfig(
            r=self.lora_config.get('rank', 16),  # Rank of adaptation
            lora_alpha=self.lora_config.get('alpha', 32),  # LoRA scaling parameter
            target_modules=self.lora_config.get('target_modules', [
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj"
            ]),
            lora_dropout=self.lora_config.get('dropout', 0.1),
            bias="none",
            task_type=TaskType.CAUSAL_LM
        )
        
        # Apply LoRA to model
        model = get_peft_model(model, lora_config)
        
        # Print trainable parameters
        model.print_trainable_parameters()
        
        return model, lora_config
    
    def calculate_memory_requirements(self, model) -> Dict[str, str]:
        """Calculate approximate memory requirements."""
        # Count parameters
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        # Estimate memory (rough approximation)
        # Model weights + gradients + optimizer states + activations
        model_memory = total_params * 2  # FP16
        training_memory = trainable_params * 8  # Gradients + optimizer states
        
        return {
            'total_parameters': f"{total_params:,}",
            'trainable_parameters': f"{trainable_params:,}",
            'trainable_percentage': f"{100 * trainable_params / total_params:.2f}%",
            'estimated_model_memory_gb': f"{model_memory * 2 / 1e9:.2f}",
            'estimated_training_memory_gb': f"{(model_memory + training_memory) * 2 / 1e9:.2f}"
        }

from transformers import BitsAndBytesConfig
import torch

class QLoRATrainer(LoRATrainer):
    """QLoRA implementation with 4-bit quantization."""
    
    def setup_quantized_model(self) -> tuple:
        """Setup model with 4-bit quantization and LoRA."""
        # 4-bit quantization config
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",  # Normal Float 4-bit
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True  # Nested quantization
        )
        
        # Load quantized model
        model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            quantization_config=bnb_config,
            device_map="auto",
            torch_dtype=torch.float16
        )
        
        # Enable gradient checkpointing for memory efficiency
        model.gradient_checkpointing_enable()
        
        # Configure LoRA for quantized model
        lora_config = LoraConfig(
            r=self.lora_config.get('rank', 16),
            lora_alpha=self.lora_config.get('alpha', 32),
            target_modules=self.lora_config.get('target_modules', [
                "q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj"
            ]),
            lora_dropout=self.lora_config.get('dropout', 0.1),
            bias="none",
            task_type=TaskType.CAUSAL_LM
        )
        
        # Apply LoRA
        model = get_peft_model(model, lora_config)
        
        return model, lora_config

from dataclasses import dataclass
from typing import Optional

@dataclass
class TrainingConfig:
    """Comprehensive training configuration."""
    
    # Model and data
    model_name: str = "microsoft/DialoGPT-medium"
    dataset_path: str = "instruction_dataset.json"
    output_dir: str = "./fine_tuned_model"
    
    # Training parameters
    learning_rate: float = 2e-4  # Higher for LoRA, lower for full fine-tuning
    batch_size: int = 4  # Adjust based on GPU memory
    gradient_accumulation_steps: int = 4  # Effective batch size = batch_size * grad_accum
    num_epochs: int = 3  # Start with fewer epochs to avoid overfitting
    max_seq_length: int = 512  # Balance between context and memory
    
    # Optimization
    optimizer: str = "adamw_torch"
    weight_decay: float = 0.01
    lr_scheduler_type: str = "cosine"
    warmup_ratio: float = 0.03  # 3% of total steps for warmup
    
    # Memory optimization
    gradient_checkpointing: bool = True
    dataloader_pin_memory: bool = True
    dataloader_num_workers: int = 4
    
    # Monitoring and saving
    eval_steps: int = 100
    save_steps: int = 500
    logging_steps: int = 10
    save_total_limit: int = 3
    
    # Early stopping
    early_stopping_patience: int = 3
    early_stopping_threshold: float = 0.01
    
    # Advanced options
    packing: bool = True  # Pack multiple samples in one sequence
    response_template: str = "### Response:"  # For response-only training

class TrainingPipeline:
    """Complete training pipeline with best practices."""
    
    def __init__(self, config: TrainingConfig):
        self.config = config
        
    def setup_training(self):
        """Setup all components for training."""
        from transformers import (
            AutoTokenizer, TrainingArguments, 
            DataCollatorForLanguageModeling
        )
        from trl import SFTTrainer
        
        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        # Setup model (LoRA or QLoRA)
        if self.config.use_qlora:
            trainer_class = QLoRATrainer(self.config.model_name, self.config.lora_config)
            model, _ = trainer_class.setup_quantized_model()
        else:
            trainer_class = LoRATrainer(self.config.model_name, self.config.lora_config)
            model, _ = trainer_class.setup_lora_model()
        
        # Training arguments
        training_args = 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,
            weight_decay=self.config.weight_decay,
            lr_scheduler_type=self.config.lr_scheduler_type,
            warmup_ratio=self.config.warmup_ratio,
            logging_steps=self.config.logging_steps,
            eval_steps=self.config.eval_steps,
            save_steps=self.config.save_steps,
            save_total_limit=self.config.save_total_limit,
            gradient_checkpointing=self.config.gradient_checkpointing,
            dataloader_pin_memory=self.config.dataloader_pin_memory,
            dataloader_num_workers=self.config.dataloader_num_workers,
            fp16=True,  # Use mixed precision
            report_to="comet_ml",  # Experiment tracking
            run_name=f"llm_twin_sft_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        )
        
        # Setup trainer
        trainer = SFTTrainer(
            model=model,
            tokenizer=tokenizer,
            train_dataset=self.load_dataset(),
            args=training_args,
            max_seq_length=self.config.max_seq_length,
            packing=self.config.packing,
            response_template=self.config.response_template
        )
        
        return trainer
    
    def train(self):
        """Execute training pipeline."""
        trainer = self.setup_training()
        
        # Start training
        trainer.train()
        
        # Save final model
        trainer.save_model()
        
        # Save training metrics
        self.save_training_metrics(trainer)
        
        return trainer
    
    def save_training_metrics(self, trainer):
        """Save training metrics and model info."""
        metrics = {
            'final_loss': trainer.state.log_history[-1].get('train_loss'),
            'total_steps': trainer.state.global_step,
            'epochs_completed': trainer.state.epoch,
            'model_config': self.config.__dict__
        }
        
        with open(f"{self.config.output_dir}/training_metrics.json", 'w') as f:
            json.dump(metrics, f, indent=2)