File size: 756 Bytes
590a501 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | """Simple model ensemble."""
from __future__ import annotations
import numpy as np
class RankEnsemble:
def __init__(self, models: list, weights: list[float] | None = None):
self.models = models
self.weights = weights or [1.0 / len(models)] * len(models)
def predict(self, X) -> np.ndarray:
preds = np.zeros(len(X))
for model, w in zip(self.models, self.weights):
if hasattr(model, "predict"):
preds += w * np.asarray(model.predict(X))
else:
import torch
model.eval()
with torch.no_grad():
xt = torch.tensor(X, dtype=torch.float32)
preds += w * model(xt).numpy()
return preds
|