CyberSecAI / models.py
invincible-jha's picture
Upload 8 files
de38b77 verified
class AIModel:
def __init__(self, model_name, model):
self.model_name = model_name
self.model = model
def analyze(self, input_data):
return self.model(input_data)
class ModelManager:
def __init__(self):
self.models = {}
def load_model(self, model_name, model):
self.models[model_name] = AIModel(model_name, model)
def update_model(self, model_name, model):
if model_name in self.models:
self.models[model_name].model = model
def enable_model(self, model_name):
if model_name in self.models:
self.models[model_name].enabled = True
def disable_model(self, model_name):
if model_name in self.models:
self.models[model_name].enabled = False
def analyze_with_model(self, model_name, input_data):
if model_name in self.models and self.models[model_name].enabled:
return self.models[model_name].analyze(input_data)
else:
return "Model not found or not enabled."
def get_model_names(self):
return list(self.models.keys())
def is_model_enabled(self, model_name):
if model_name in self.models:
return self.models[model_name].enabled
else:
return False