#!/usr/bin/env python3
"""
train_multilabel.py - Multi-label classifier training for 5-tag system
Uses transformers for efficient training on CPU
"""

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer, get_linear_schedule_with_warmup
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import f1_score
import json
from pathlib import Path

# Configuration
MODEL_NAME = "distilbert-base-uncased"
MAX_LENGTH = 128
BATCH_SIZE = 16
LEARNING_RATE = 2e-5
NUM_EPOCHS = 10
THRESHOLD = 0.40
DEVICE = torch.device("cpu")

class MessageDataset(Dataset):
    """Dataset for multi-label message classification"""

    def __init__(self, data_path: Path, tokenizer, tag2id: dict, max_length: int = 128):
        self.tokenizer = tokenizer
        self.tag2id = tag2id
        self.max_length = max_length
        self.messages = []
        self.labels = []

        # Load data
        with open(data_path, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                self.messages.append(item['message'])

                # Create multi-label binary vector
                label_vector = [0] * len(tag2id)
                for tag in item.get('tags', []):
                    if tag in tag2id:
                        label_vector[tag2id[tag]] = 1
                self.labels.append(label_vector)

    def __len__(self):
        return len(self.messages)

    def __getitem__(self, idx):
        message = self.messages[idx]
        labels = self.labels[idx]

        # Tokenize
        encoding = self.tokenizer(
            message,
            truncation=True,
            padding='max_length',
            max_length=self.max_length,
            return_tensors='pt'
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(labels, dtype=torch.float)
        }

class MultiLabelClassifier(nn.Module):
    """Multi-label classifier model"""

    def __init__(self, model_name: str, num_labels: int):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        # Use [CLS] token representation or pooler_output if available
        if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None:
            pooled_output = outputs.pooler_output
        else:
            # For DistilBERT - use [CLS] token (first token)
            pooled_output = outputs.last_hidden_state[:, 0, :]

        output = self.dropout(pooled_output)
        logits = self.classifier(output)
        return logits

def train_epoch(model, dataloader, optimizer, scheduler, criterion):
    """Train for one epoch"""
    model.train()
    total_loss = 0

    for batch in dataloader:
        optimizer.zero_grad()

        input_ids = batch['input_ids'].to(DEVICE)
        attention_mask = batch['attention_mask'].to(DEVICE)
        labels = batch['labels'].to(DEVICE)

        logits = model(input_ids, attention_mask)
        loss = criterion(logits, labels)

        loss.backward()
        optimizer.step()
        scheduler.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

def evaluate(model, dataloader, criterion, threshold=0.5):
    """Evaluate model on validation set"""
    model.eval()
    all_predictions = []
    all_labels = []

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch['input_ids'].to(DEVICE)
            attention_mask = batch['attention_mask'].to(DEVICE)
            labels = batch['labels'].to(DEVICE)

            logits = model(input_ids, attention_mask)

            # Apply sigmoid and threshold
            predictions = torch.sigmoid(logits)
            predictions = (predictions > threshold).float()

            all_predictions.extend(predictions.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    # Calculate metrics
    all_predictions = np.array(all_predictions)
    all_labels = np.array(all_labels)

    # Micro-averaged F1 (as specified in requirements)
    micro_f1 = f1_score(all_labels, all_predictions, average='micro', zero_division=0)
    macro_f1 = f1_score(all_labels, all_predictions, average='macro', zero_division=0)

    return micro_f1, macro_f1

# Training execution
print(f"Starting training...")
print(f"Model: {MODEL_NAME}")
print(f"Device: {DEVICE}")

# Initialize tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = MultiLabelClassifier(MODEL_NAME, num_labels=5)
model.to(DEVICE)

# Setup training
criterion = nn.BCEWithLogitsLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

# Training loop
for epoch in range(NUM_EPOCHS):
    train_loss = train_epoch(model, train_loader, optimizer, scheduler, criterion)
    val_micro_f1, val_macro_f1 = evaluate(model, val_loader, criterion, THRESHOLD)

    print(f"Epoch {epoch + 1}/{NUM_EPOCHS}")
    print(f"Train Loss: {train_loss:.4f}")
    print(f"Val Micro-F1: {val_micro_f1:.4f}")
    print(f"Val Macro-F1: {val_macro_f1:.4f}")

print("Training complete!")

#!/usr/bin/env python3
"""
serve.py - FastAPI service for message and conversation classification
Provides REST API endpoints for real-time classification
"""

import json
import time
from pathlib import Path
from typing import Dict, List, Optional
from collections import deque, defaultdict
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
import numpy as np

# Initialize FastAPI app
app = FastAPI(
    title="5-Tag Message Classifier API",
    description="Multi-label message and conversation classifier for UnifiedBeez",
    version="1.0.0"
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables for model
model = None
tokenizer = None
config = None
tag2id = None
id2tag = None
conversation_store = defaultdict(lambda: deque(maxlen=3))

# Request/Response models
class MessageRequest(BaseModel):
    message: str

class ConversationRequest(BaseModel):
    conversation_id: str
    message: str

class TagScore(BaseModel):
    tag: str
    confidence: float

class MessageResponse(BaseModel):
    threshold: float
    selected: List[TagScore]
    all: List[TagScore]

class ConversationResponse(BaseModel):
    conversation_id: str
    threshold: float
    conversation_tag: Optional[TagScore]
    switched: bool
    evidence: List[Dict]

# Model architecture (must match training)
class MultiLabelClassifier(nn.Module):
    """Multi-label classifier model"""

    def __init__(self, model_name: str, num_labels: int):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(self.bert.config.hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        # Use [CLS] token representation or pooler_output if available
        if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None:
            pooled_output = outputs.pooler_output
        else:
            # For DistilBERT - use [CLS] token (first token)
            pooled_output = outputs.last_hidden_state[:, 0, :]
        output = self.dropout(pooled_output)
        logits = self.classifier(output)
        return logits

def load_model_on_startup():
    """Load model and configuration at startup"""
    global model, tokenizer, config, tag2id, id2tag

    base_dir = Path(__file__).parent
    model_dir = base_dir / "model_out"

    if not model_dir.exists():
        raise RuntimeError("Model directory not found. Please train the model first.")

    print("Loading model...")

    # Load config
    with open(model_dir / "config.runtime.json", 'r') as f:
        config = json.load(f)

    # Load tag2id
    with open(model_dir / "tag2id.json", 'r') as f:
        tag2id = json.load(f)
    id2tag = {v: k for k, v in tag2id.items()}

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    # Load model
    model = MultiLabelClassifier(config['model_name'], config['num_labels'])
    model.load_state_dict(torch.load(model_dir / "model.pt", map_location='cpu'))
    model.eval()

    print(f"Model loaded successfully!")
    print(f"Tags: {list(tag2id.keys())}")
    print(f"Threshold: {config['threshold']}")

def classify_message_internal(message: str) -> Dict:
    """Internal function to classify a message"""
    # Tokenize
    inputs = tokenizer(
        message,
        truncation=True,
        padding='max_length',
        max_length=config['max_length'],
        return_tensors='pt'
    )

    # Predict
    with torch.no_grad():
        logits = model(inputs['input_ids'], inputs['attention_mask'])
        probabilities = torch.sigmoid(logits).squeeze().numpy()

    # Get predictions
    threshold = config['threshold']
    selected = []
    all_scores = []

    for i, prob in enumerate(probabilities):
        tag = id2tag[i]
        score = {"tag": tag, "confidence": float(prob)}
        all_scores.append(score)
        if prob > threshold:
            selected.append(score)

    # Sort by confidence
    selected.sort(key=lambda x: x['confidence'], reverse=True)
    all_scores.sort(key=lambda x: x['confidence'], reverse=True)

    return {
        "selected": selected,
        "all": all_scores,
        "probabilities": probabilities
    }

@app.on_event("startup")
async def startup_event():
    """Load model on startup"""
    load_model_on_startup()
    print("API ready to serve requests")

@app.get("/")
async def root():
    """Root endpoint with API information"""
    return {
        "service": "5-Tag Message Classifier",
        "version": "1.0.0",
        "endpoints": [
            "/classify-message",
            "/classify-conversation",
            "/health"
        ],
        "tags": list(tag2id.keys()) if tag2id else [],
        "model_loaded": model is not None
    }

@app.get("/health")
async def health():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "model_loaded": model is not None
    }

@app.post("/classify-message", response_model=MessageResponse)
async def classify_message(request: MessageRequest):
    """
    Classify a single message into one or more tags.
    Returns selected tags (above threshold) and all tag scores.
    """
    if not model:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        start_time = time.time()
        result = classify_message_internal(request.message)
        latency_ms = (time.time() - start_time) * 1000

        response = MessageResponse(
            threshold=config['threshold'],
            selected=result['selected'],
            all=result['all']
        )

        # Log for monitoring
        print(f"Message classified in {latency_ms:.1f}ms - Tags: {[t['tag'] for t in result['selected']]}")

        return response

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/classify-conversation", response_model=ConversationResponse)
async def classify_conversation(request: ConversationRequest):
    """
    Classify a message and update conversation-level tag.
    Maintains state over last 3 messages and detects tag switches.
    """
    if not model:
        raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        # Classify the new message
        result = classify_message_internal(request.message)

        # Store message with its scores
        message_data = {
            "message": request.message,
            "probabilities": result['probabilities'],
            "top": result['selected'][:1] if result['selected'] else []
        }

        # Get previous conversation tag (if exists)
        conversation_history = conversation_store[request.conversation_id]
        previous_tag = None

        if conversation_history:
            # Calculate previous conversation tag
            avg_probs = np.mean([msg['probabilities'] for msg in conversation_history], axis=0)
            max_idx = np.argmax(avg_probs)
            if avg_probs[max_idx] >= config['threshold']:
                previous_tag = id2tag[max_idx]

        # Add new message to history
        conversation_history.append(message_data)

        # Calculate new conversation tag (average of last 3 messages)
        recent_messages = list(conversation_history)
        avg_probabilities = np.mean([msg['probabilities'] for msg in recent_messages], axis=0)

        # Find the top tag
        max_idx = np.argmax(avg_probabilities)
        max_prob = avg_probabilities[max_idx]

        conversation_tag = None
        if max_prob >= config['threshold']:
            conversation_tag = TagScore(
                tag=id2tag[max_idx],
                confidence=float(max_prob)
            )

        # Check if tag switched
        current_tag = conversation_tag.tag if conversation_tag else None
        switched = (current_tag != previous_tag) if previous_tag else False

        # Prepare evidence (last messages)
        evidence = []
        for msg in recent_messages[-2:]:
            evidence.append({
                "message": msg['message'][:50] + "..." if len(msg['message']) > 50 else msg['message'],
                "top": msg['top']
            })

        response = ConversationResponse(
            conversation_id=request.conversation_id,
            threshold=config['threshold'],
            conversation_tag=conversation_tag,
            switched=switched,
            evidence=evidence
        )

        # Log for monitoring
        if switched:
            print(f"Conversation {request.conversation_id}: Tag switched from {previous_tag} to {current_tag}")

        return response

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/conversation/{conversation_id}")
async def clear_conversation(conversation_id: str):
    """Clear conversation history (useful for testing)"""
    if conversation_id in conversation_store:
        del conversation_store[conversation_id]
        return {"message": f"Conversation {conversation_id} cleared"}
    return {"message": f"Conversation {conversation_id} not found"}

def main():
    """Run the FastAPI server"""
    print("Starting 5-Tag Message Classifier API...")
    print("API will be available at http://localhost:8000")
    print("Documentation at http://localhost:8000/docs")

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info"
    )

if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
prepare_tag_dataset.py - Prepare training data for multi-label classification
Handles deduplication, train/val/test splitting, and tag mapping
"""

import json
from pathlib import Path
from sklearn.model_selection import train_test_split

def prepare_dataset():
    """Prepare dataset from seed data"""
    base_dir = Path(__file__).parent.parent
    data_dir = base_dir / "data"
    output_dir = base_dir / "prepared_dataset"
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load seed data
    seed_file = data_dir / "seed.jsonl"
    data = []

    with open(seed_file, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:  # Skip blank lines
                try:
                    item = json.loads(line)
                    data.append(item)
                except json.JSONDecodeError:
                    continue

    print(f"Loaded {len(data)} valid messages")

    # Deduplication
    seen_messages = set()
    unique_data = []

    for item in data:
        if item['message'] not in seen_messages:
            unique_data.append(item)
            seen_messages.add(item['message'])

    print(f"After deduplication: {len(unique_data)} messages")

    # Extract all tags
    all_tags = set()
    for item in unique_data:
        all_tags.update(item.get('tags', []))

    # Create tag2id mapping
    tag2id = {tag: idx for idx, tag in enumerate(sorted(all_tags))}

    # Stratified split (67% train, 15% val, 18% test)
    train_data, temp = train_test_split(unique_data, test_size=0.33, random_state=42)
    val_data, test_data = train_test_split(temp, test_size=0.53, random_state=42)

    print(f"\nDataset split:")
    print(f"  Train: {len(train_data)}")
    print(f"  Val: {len(val_data)}")
    print(f"  Test: {len(test_data)}")

    # Save splits
    for name, split_data in [('train', train_data), ('val', val_data), ('test', test_data)]:
        output_file = output_dir / f"{name}.jsonl"
        with open(output_file, 'w', encoding='utf-8') as f:
            for item in split_data:
                f.write(json.dumps(item) + '\n')

    # Save tag2id
    with open(output_dir / "tag2id.json", 'w') as f:
        json.dump(tag2id, f, indent=2)

    # Generate statistics
    stats = {
        "total_messages": len(unique_data),
        "num_tags": len(all_tags),
        "tags": list(tag2id.keys()),
        "train_size": len(train_data),
        "val_size": len(val_data),
        "test_size": len(test_data)
    }

    with open(output_dir / "stats.json", 'w') as f:
        json.dump(stats, f, indent=2)

    print("\nData preparation complete!")
    print(f"Output directory: {output_dir}")

if __name__ == "__main__":
    prepare_dataset()