quant_test / models /ensemble.py
lucky-loster's picture
Upload folder using huggingface_hub
590a501 verified
Raw
History Blame Contribute Delete
756 Bytes
"""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