Files
LLM_Engineering_OLD/week8/community_contributions/ensemble-updated/ensemble_agent.py

52 lines
2.0 KiB
Python

import pandas as pd
from sklearn.linear_model import LinearRegression
import joblib
from agents.agent import Agent
from agents.specialist_agent import SpecialistAgent
from agents.frontier_agent import FrontierAgent
from agents.random_forest_agent import RandomForestAgent
from agents.xg_boost_agent import XGBoostAgent
class EnsembleAgent(Agent):
name = "Ensemble Agent"
color = Agent.YELLOW
def __init__(self, collection):
"""
Create an instance of Ensemble, by creating each of the models
And loading the weights of the Ensemble
"""
self.log("Initializing Ensemble Agent")
self.specialist = SpecialistAgent()
self.frontier = FrontierAgent(collection)
self.random_forest = RandomForestAgent()
self.xg_boost = XGBoostAgent()
self.model = joblib.load('ensemble_model.pkl')
self.log("Ensemble Agent is ready")
def price(self, description: str) -> float:
"""
Run this ensemble model
Ask each of the models to price the product
Then use the Linear Regression model to return the weighted price
:param description: the description of a product
:return: an estimate of its price
"""
self.log("Running Ensemble Agent - collaborating with specialist, frontier, xg boost and random forest agents")
specialist = self.specialist.price(description)
frontier = self.frontier.price(description)
random_forest = self.random_forest.price(description)
xg_boost = self.xg_boost.price(description)
X = pd.DataFrame({
'Specialist': [specialist],
'Frontier': [frontier],
'RandomForest': [random_forest],
'XGBoost' : [xg_boost],
'Min': [min(specialist, frontier, random_forest, xg_boost)],
'Max': [max(specialist, frontier, random_forest, xg_boost)],
})
y = max(0, self.model.predict(X)[0])
self.log(f"Ensemble Agent complete - returning ${y:.2f}")
return y