Capstone / _hf_deploy /agents /ensemble_agent.py
prabalGaur's picture
Deploy updated app and community contributions
ed9ecbf verified
Raw
History Blame Contribute Delete
2.15 kB
from agents.agent import Agent
from agents.specialist_agent import SpecialistAgent
from agents.frontier_agent import FrontierAgent
from agents.neural_network_agent import NeuralNetworkAgent
from agents.preprocessor import Preprocessor
from free_config import WEIGHTS_WITH_SPECIALIST, WEIGHTS_WITHOUT_SPECIALIST
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.neural_network = NeuralNetworkAgent()
self.preprocessor = Preprocessor()
self.use_specialist = self.specialist.available
if self.use_specialist:
mode = getattr(self.specialist, "mode", "on")
self.log(f"Ensemble Agent is ready (Frontier + Specialist [{mode}] + Neural Network)")
else:
self.log("Ensemble Agent is ready (Frontier + Neural Network only)")
def price(self, description: str) -> float:
"""
Run this ensemble model
Ask each of the models to price the product
Then return the weighted average price
"""
self.log("Running Ensemble Agent - preprocessing text")
rewrite = self.preprocessor.preprocess(description)
self.log(f"Pre-processed text using {self.preprocessor.model_name}")
frontier = self.frontier.price(rewrite)
neural_network = self.neural_network.price(rewrite)
if self.use_specialist:
specialist = self.specialist.price(rewrite)
wf, ws, wn = WEIGHTS_WITH_SPECIALIST
combined = frontier * wf + specialist * ws + neural_network * wn
else:
wf, wn = WEIGHTS_WITHOUT_SPECIALIST
combined = frontier * wf + neural_network * wn
self.log(f"Ensemble Agent complete - returning ${combined:.2f}")
return combined