File size: 2,414 Bytes
cff2458 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import argparse
import os
import importlib
def check_plugins(loaded_plugins):
print("Loaded plugins:")
for plugin in loaded_plugins:
print(f"- {plugin}")
def train_model(dataset_name, plugins):
dataset = {'train': []} # Placeholder for training data
model = "FlowModel"
for plugin in plugins:
if hasattr(plugin, 'modify_model'):
model = plugin.modify_model(model)
for plugin in plugins:
if hasattr(plugin, 'on_train_start'):
plugin.on_train_start()
print(f"Training started on dataset: {dataset_name}")
for plugin in plugins:
if hasattr(plugin, 'on_train_end'):
plugin.on_train_end()
print("Training finished.")
def load_plugins():
plugins_dir = './plugins'
plugins = []
if not os.path.exists(plugins_dir):
os.makedirs(plugins_dir)
print(f"Plugins directory created at {plugins_dir}. Add your plugins there!")
for filename in os.listdir(plugins_dir):
if filename.endswith('.py') and filename != '__init__.py':
plugin_name = filename[:-3]
try:
plugin_module = importlib.import_module(f'plugins.{plugin_name}')
plugin_class = getattr(plugin_module, plugin_name.title().replace('_', ''), None)
if plugin_class:
plugins.append(plugin_class())
print(f"Plugin {plugin_name} loaded.")
else:
print(f"No class found in plugin {plugin_name}.")
except Exception as e:
print(f"Failed to load plugin {plugin_name}: {e}")
return plugins
def predict_model(plugins):
print("Prediction started.")
for plugin in plugins:
if hasattr(plugin, 'on_predict'):
plugin.on_predict()
print("Prediction finished.")
def main():
parser = argparse.ArgumentParser(description="FlowModel CLI")
parser.add_argument('command', choices=['train', 'predict', 'check_plugins'], help="Command to run")
args = parser.parse_args()
plugins, loaded_plugins = load_plugins()
if args.command == 'train':
plugins = load_plugins()
train_model("mnist", plugins)
elif args.command == 'predict':
predict_model(plugins)
elif args.command == 'check_plugins':
check_plugins(loaded_plugins)
if __name__ == "__main__":
main()
|