Upload runpod_handler.py

runpod_handler.py

@@ -0,0 +1,164 @@
+import runpod
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoModel
+import os
+import logging
+from typing import Dict, Any, List
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class MultiOutputClassifier(nn.Module):
+    """Multi-output classifier"""
+    def __init__(self, encoder, hidden_size, num_classes=3, num_levels=3):
+        super().__init__()
+        self.encoder = encoder
+        self.classifiers = nn.ModuleList([
+            nn.Linear(hidden_size, num_levels) for _ in range(num_classes)
+        ])
+    
+    def forward(self, input_ids, attention_mask):
+        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_states = outputs.last_hidden_state
+        
+        attention_mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
+        sum_hidden = torch.sum(hidden_states * attention_mask_expanded, dim=1)
+        sum_mask = torch.clamp(attention_mask_expanded.sum(dim=1), min=1e-9)
+        pooled_output = sum_hidden / sum_mask
+        
+        logits = [classifier(pooled_output) for classifier in self.classifiers]
+        logits = torch.stack(logits, dim=1)
+        
+        return logits
+
+# Global model instance
+model = None
+tokenizer = None
+device = None
+
+def load_model():
+    """Load model once at startup"""
+    global model, tokenizer, device
+    
+    model_path = "/app/model"
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    logger.info(f"Loading model from {model_path}")
+    logger.info(f"Using device: {device}")
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load encoder
+    encoder = AutoModel.from_pretrained(
+        model_path,
+        use_safetensors=True,
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    )
+    
+    # Initialize model
+    hidden_size = encoder.config.hidden_size
+    model = MultiOutputClassifier(
+        encoder=encoder,
+        hidden_size=hidden_size,
+        num_classes=3,
+        num_levels=3
+    )
+    
+    # Load classification heads
+    classifier_heads_path = os.path.join(model_path, "classifier_heads.pt")
+    if os.path.exists(classifier_heads_path):
+        logger.info(f"Loading classification heads")
+        checkpoint = torch.load(classifier_heads_path, map_location=device)
+        classifiers_list = checkpoint['classifiers']
+        
+        for i, classifier in enumerate(model.classifiers):
+            classifier.load_state_dict(classifiers_list[i])
+    
+    model.to(device)
+    model.eval()
+    logger.info("✓ Model loaded and ready!")
+    
+    return model
+
+def handler(job: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    RunPod serverless handler
+    
+    Input format:
+    {
+        "input": {
+            "text": "Your text here" OR ["text1", "text2"],
+            "max_length": 128,  # optional
+            "return_scores": true  # optional
+        }
+    }
+    """
+    try:
+        job_input = job["input"]
+        
+        # Extract inputs
+        text_input = job_input.get("text", job_input.get("inputs", ""))
+        max_length = job_input.get("max_length", 128)
+        return_scores = job_input.get("return_scores", True)
+        
+        # Handle both single string and list
+        if isinstance(text_input, str):
+            texts = [text_input]
+        else:
+            texts = text_input
+        
+        # Tokenize
+        encoded = tokenizer(
+            texts,
+            truncation=True,
+            padding='max_length',
+            max_length=max_length,
+            return_tensors='pt'
+        )
+        
+        input_ids = encoded['input_ids'].to(device)
+        attention_mask = encoded['attention_mask'].to(device)
+        
+        # Inference
+        with torch.no_grad():
+            logits = model(input_ids=input_ids, attention_mask=attention_mask)
+            probs = torch.softmax(logits, dim=2)
+            preds = torch.argmax(logits, dim=2)
+        
+        # Format results
+        class_names = ['vyhruzky', 'vulgarity', 'rasismus']
+        level_labels = {0: 'none', 1: 'moderate', 2: 'severe'}
+        
+        results = []
+        for i in range(len(texts)):
+            result = {}
+            for j, class_name in enumerate(class_names):
+                pred_class = preds[i, j].item()
+                pred_prob = probs[i, j, pred_class].item()
+                
+                result[class_name] = {
+                    "label": level_labels[pred_class],
+                    "level": pred_class
+                }
+                
+                if return_scores:
+                    result[class_name]["score"] = round(pred_prob, 4)
+            
+            results.append(result)
+        
+        return {"output": results}
+    
+    except Exception as e:
+        logger.error(f"Error in handler: {str(e)}")
+        return {"error": str(e)}
+
+if __name__ == "__main__":
+    logger.info("Starting RunPod serverless handler...")
+    load_model()
+    logger.info("Starting RunPod serverless worker...")
+    runpod.serverless.start({"handler": handler})