From 0a89751f84b722098fb46c0254777191d860fe39 Mon Sep 17 00:00:00 2001
From: Pop John <john.pop@airgap.io>
Date: Wed, 7 Aug 2024 15:48:28 +0300
Subject: [PATCH] bug fixes

---
 .../ms-panel-model-training.component.html             | 10 +++++-----
 .../ms-panel-model/ms-panel-model.component.html       |  8 ++++----
 .../ms-panel-model/ms-panel-model.component.ts         |  5 +++--
 .../app/styles/theme/components/_ms-components.scss    |  1 +
 .../src/app/styles/theme/components/_ms-select.scss    |  7 +++++++
 machine_learning_core/multiflow/prune.py               |  2 ++
 6 files changed, 22 insertions(+), 11 deletions(-)
 create mode 100644 frontend/src/app/styles/theme/components/_ms-select.scss

diff --git a/frontend/src/app/modules/shared/components/ms-panel-model-training/ms-panel-model-training.component.html b/frontend/src/app/modules/shared/components/ms-panel-model-training/ms-panel-model-training.component.html
index a210dc96..461955dc 100644
--- a/frontend/src/app/modules/shared/components/ms-panel-model-training/ms-panel-model-training.component.html
+++ b/frontend/src/app/modules/shared/components/ms-panel-model-training/ms-panel-model-training.component.html
@@ -22,13 +22,13 @@
 	<div class="form-field-container" [formGroupName]="controlKey">
 		<mat-form-field appearance="outline" subscriptSizing="dynamic">
 			<mat-select formControlName="model" placeholder="Select a model...">
-				<mat-optgroup label="Predefined Models">
-					<mat-option>
-						<ngx-mat-select-search [formControl]="searchModel"> </ngx-mat-select-search>
-					</mat-option>
+				<mat-option>
+					<ngx-mat-select-search [formControl]="searchModel"> </ngx-mat-select-search>
+				</mat-option>
 
+				<mat-optgroup label="Predefined Models">
 					@if (filteredModels | async; as models) { @for (model of models; track model.name) {
-					<mat-option [value]="model.name">
+					<mat-option [value]="model.name" [disabled]="!model.isTrained">
 						{{ model.name }}
 						@if (!model.isTrained) {
 						<mat-icon class="mat-error" fontSet="ms" fontIcon="icon-X"></mat-icon>
diff --git a/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.html b/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.html
index 90211a12..3a535acc 100644
--- a/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.html
+++ b/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.html
@@ -31,11 +31,11 @@
 	<div class="form-field-container" [formGroupName]="controlKey">
 		<mat-form-field appearance="outline" subscriptSizing="dynamic">
 			<mat-select formControlName="model" placeholder="Select a model...">
-				<mat-optgroup label="Predefined Models">
-					<mat-option>
-						<ngx-mat-select-search [formControl]="searchModel"> </ngx-mat-select-search>
-					</mat-option>
+				<mat-option>
+					<ngx-mat-select-search [formControl]="searchModel"> </ngx-mat-select-search>
+				</mat-option>
 
+				<mat-optgroup label="Predefined Models">
 					@if (filteredModels | async; as models) { @for (model of models; track model.name) {
 					<mat-option [value]="model.name" [disabled]="!model.isTrained">
 						{{ model.name }}
diff --git a/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.ts b/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.ts
index 438907e2..dfc142d5 100644
--- a/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.ts
+++ b/frontend/src/app/modules/shared/components/ms-panel-model/ms-panel-model.component.ts
@@ -16,7 +16,7 @@
 
 import { CommonModule } from '@angular/common';
 import { Component, Input, OnChanges, OnDestroy, OnInit, SimpleChanges, inject } from '@angular/core';
-import { ControlContainer, FormControl, FormGroup, ReactiveFormsModule, Validators } from '@angular/forms';
+import { ControlContainer, FormControl, FormGroup, FormsModule, ReactiveFormsModule, Validators } from '@angular/forms';
 import { MatCardModule } from '@angular/material/card';
 import { MatFormFieldModule } from '@angular/material/form-field';
 import { MatIconModule } from '@angular/material/icon';
@@ -44,12 +44,13 @@ import { MsSpiningIndicatorComponent } from '../ms-spining-indicator/ms-spining-
 	selector: 'ms-panel-model',
 	standalone: true,
 	imports: [
+		CommonModule,
+		FormsModule,
 		ReactiveFormsModule,
 		MatCardModule,
 		MatFormFieldModule,
 		MatSelectModule,
 		MatIconModule,
-		CommonModule,
 		NgxMatSelectSearchModule,
 		RouterLink,
 		MatProgressSpinnerModule,
diff --git a/frontend/src/app/styles/theme/components/_ms-components.scss b/frontend/src/app/styles/theme/components/_ms-components.scss
index e2b019c5..e2bdf965 100644
--- a/frontend/src/app/styles/theme/components/_ms-components.scss
+++ b/frontend/src/app/styles/theme/components/_ms-components.scss
@@ -26,3 +26,4 @@
 @forward './ms-tooltip';
 @forward './ms-snackbar';
 @forward './ms-drawer';
+@forward './ms-select';
diff --git a/frontend/src/app/styles/theme/components/_ms-select.scss b/frontend/src/app/styles/theme/components/_ms-select.scss
new file mode 100644
index 00000000..0116b738
--- /dev/null
+++ b/frontend/src/app/styles/theme/components/_ms-select.scss
@@ -0,0 +1,7 @@
+.mdc-list-item__primary-text {
+	width: 100%;
+}
+
+.mat-mdc-optgroup-label {
+	padding: 5px 16px;
+}
diff --git a/machine_learning_core/multiflow/prune.py b/machine_learning_core/multiflow/prune.py
index e5fea05d..55166e14 100644
--- a/machine_learning_core/multiflow/prune.py
+++ b/machine_learning_core/multiflow/prune.py
@@ -209,6 +209,8 @@ def main(args):
     runtimes = pd.DataFrame(runtimes)
     runtimes.to_csv(os.path.join(args.output_dir, f"{args.model}_{args.pruner}_seed{args.seed}_runtimes.csv"), index=False)
 
+    print(f"Pruning task completed successfully! The pruned weights are saved in the directory: {args.output_dir}")
+
 if __name__ == "__main__":
     parser = ArgumentParser()
     parser.add_argument('-p', '--pruner', type=str, required=True, choices=available_pruners)