├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── LICENSE
├── NOTICE
├── README.md
├── THIRD-PARTY-LICENSES
├── anollm
├── __init__.py
├── anollm.py
├── anollm_dataset.py
├── anollm_trainer.py
└── anollm_utils.py
├── evaluate_anollm.py
├── evaluate_baselines.py
├── figs
└── overview.png
├── requirements.txt
├── scripts
├── exp1-mixed_benchmark
│ ├── run_anollm.sh
│ └── run_baselines.sh
├── exp2-odds
│ ├── run_anollm.sh
│ └── run_baselines.sh
├── exp3-binning_effect
│ └── run_binning_odds.sh
└── exp4-model_size
│ ├── run_anollm_1.7B_mixed.sh
│ └── run_anollm_1.7B_odds.sh
├── src
├── __init__.py
├── baselines
│ ├── dte.py
│ └── icl.py
├── data_utils.py
├── get_avg_results.py
└── get_results.py
└── train_anollm.py
/CODE_OF_CONDUCT.md:
--------------------------------------------------------------------------------
1 | ## Code of Conduct
2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
4 | opensource-codeofconduct@amazon.com with any additional questions or comments.
5 |
--------------------------------------------------------------------------------
/CONTRIBUTING.md:
--------------------------------------------------------------------------------
1 | # Contributing Guidelines
2 |
3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional
4 | documentation, we greatly value feedback and contributions from our community.
5 |
6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary
7 | information to effectively respond to your bug report or contribution.
8 |
9 |
10 | ## Reporting Bugs/Feature Requests
11 |
12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features.
13 |
14 | When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already
15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful:
16 |
17 | * A reproducible test case or series of steps
18 | * The version of our code being used
19 | * Any modifications you've made relevant to the bug
20 | * Anything unusual about your environment or deployment
21 |
22 |
23 | ## Contributing via Pull Requests
24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that:
25 |
26 | 1. You are working against the latest source on the *main* branch.
27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already.
28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted.
29 |
30 | To send us a pull request, please:
31 |
32 | 1. Fork the repository.
33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change.
34 | 3. Ensure local tests pass.
35 | 4. Commit to your fork using clear commit messages.
36 | 5. Send us a pull request, answering any default questions in the pull request interface.
37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation.
38 |
39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and
40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/).
41 |
42 |
43 | ## Finding contributions to work on
44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' issues is a great place to start.
45 |
46 |
47 | ## Code of Conduct
48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
50 | opensource-codeofconduct@amazon.com with any additional questions or comments.
51 |
52 |
53 | ## Security issue notifications
54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue.
55 |
56 |
57 | ## Licensing
58 |
59 | See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution.
60 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 |
2 | Apache License
3 | Version 2.0, January 2004
4 | http://www.apache.org/licenses/
5 |
6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
7 |
8 | 1. Definitions.
9 |
10 | "License" shall mean the terms and conditions for use, reproduction,
11 | and distribution as defined by Sections 1 through 9 of this document.
12 |
13 | "Licensor" shall mean the copyright owner or entity authorized by
14 | the copyright owner that is granting the License.
15 |
16 | "Legal Entity" shall mean the union of the acting entity and all
17 | other entities that control, are controlled by, or are under common
18 | control with that entity. For the purposes of this definition,
19 | "control" means (i) the power, direct or indirect, to cause the
20 | direction or management of such entity, whether by contract or
21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the
22 | outstanding shares, or (iii) beneficial ownership of such entity.
23 |
24 | "You" (or "Your") shall mean an individual or Legal Entity
25 | exercising permissions granted by this License.
26 |
27 | "Source" form shall mean the preferred form for making modifications,
28 | including but not limited to software source code, documentation
29 | source, and configuration files.
30 |
31 | "Object" form shall mean any form resulting from mechanical
32 | transformation or translation of a Source form, including but
33 | not limited to compiled object code, generated documentation,
34 | and conversions to other media types.
35 |
36 | "Work" shall mean the work of authorship, whether in Source or
37 | Object form, made available under the License, as indicated by a
38 | copyright notice that is included in or attached to the work
39 | (an example is provided in the Appendix below).
40 |
41 | "Derivative Works" shall mean any work, whether in Source or Object
42 | form, that is based on (or derived from) the Work and for which the
43 | editorial revisions, annotations, elaborations, or other modifications
44 | represent, as a whole, an original work of authorship. For the purposes
45 | of this License, Derivative Works shall not include works that remain
46 | separable from, or merely link (or bind by name) to the interfaces of,
47 | the Work and Derivative Works thereof.
48 |
49 | "Contribution" shall mean any work of authorship, including
50 | the original version of the Work and any modifications or additions
51 | to that Work or Derivative Works thereof, that is intentionally
52 | submitted to Licensor for inclusion in the Work by the copyright owner
53 | or by an individual or Legal Entity authorized to submit on behalf of
54 | the copyright owner. For the purposes of this definition, "submitted"
55 | means any form of electronic, verbal, or written communication sent
56 | to the Licensor or its representatives, including but not limited to
57 | communication on electronic mailing lists, source code control systems,
58 | and issue tracking systems that are managed by, or on behalf of, the
59 | Licensor for the purpose of discussing and improving the Work, but
60 | excluding communication that is conspicuously marked or otherwise
61 | designated in writing by the copyright owner as "Not a Contribution."
62 |
63 | "Contributor" shall mean Licensor and any individual or Legal Entity
64 | on behalf of whom a Contribution has been received by Licensor and
65 | subsequently incorporated within the Work.
66 |
67 | 2. Grant of Copyright License. Subject to the terms and conditions of
68 | this License, each Contributor hereby grants to You a perpetual,
69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
70 | copyright license to reproduce, prepare Derivative Works of,
71 | publicly display, publicly perform, sublicense, and distribute the
72 | Work and such Derivative Works in Source or Object form.
73 |
74 | 3. Grant of Patent License. Subject to the terms and conditions of
75 | this License, each Contributor hereby grants to You a perpetual,
76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
77 | (except as stated in this section) patent license to make, have made,
78 | use, offer to sell, sell, import, and otherwise transfer the Work,
79 | where such license applies only to those patent claims licensable
80 | by such Contributor that are necessarily infringed by their
81 | Contribution(s) alone or by combination of their Contribution(s)
82 | with the Work to which such Contribution(s) was submitted. If You
83 | institute patent litigation against any entity (including a
84 | cross-claim or counterclaim in a lawsuit) alleging that the Work
85 | or a Contribution incorporated within the Work constitutes direct
86 | or contributory patent infringement, then any patent licenses
87 | granted to You under this License for that Work shall terminate
88 | as of the date such litigation is filed.
89 |
90 | 4. Redistribution. You may reproduce and distribute copies of the
91 | Work or Derivative Works thereof in any medium, with or without
92 | modifications, and in Source or Object form, provided that You
93 | meet the following conditions:
94 |
95 | (a) You must give any other recipients of the Work or
96 | Derivative Works a copy of this License; and
97 |
98 | (b) You must cause any modified files to carry prominent notices
99 | stating that You changed the files; and
100 |
101 | (c) You must retain, in the Source form of any Derivative Works
102 | that You distribute, all copyright, patent, trademark, and
103 | attribution notices from the Source form of the Work,
104 | excluding those notices that do not pertain to any part of
105 | the Derivative Works; and
106 |
107 | (d) If the Work includes a "NOTICE" text file as part of its
108 | distribution, then any Derivative Works that You distribute must
109 | include a readable copy of the attribution notices contained
110 | within such NOTICE file, excluding those notices that do not
111 | pertain to any part of the Derivative Works, in at least one
112 | of the following places: within a NOTICE text file distributed
113 | as part of the Derivative Works; within the Source form or
114 | documentation, if provided along with the Derivative Works; or,
115 | within a display generated by the Derivative Works, if and
116 | wherever such third-party notices normally appear. The contents
117 | of the NOTICE file are for informational purposes only and
118 | do not modify the License. You may add Your own attribution
119 | notices within Derivative Works that You distribute, alongside
120 | or as an addendum to the NOTICE text from the Work, provided
121 | that such additional attribution notices cannot be construed
122 | as modifying the License.
123 |
124 | You may add Your own copyright statement to Your modifications and
125 | may provide additional or different license terms and conditions
126 | for use, reproduction, or distribution of Your modifications, or
127 | for any such Derivative Works as a whole, provided Your use,
128 | reproduction, and distribution of the Work otherwise complies with
129 | the conditions stated in this License.
130 |
131 | 5. Submission of Contributions. Unless You explicitly state otherwise,
132 | any Contribution intentionally submitted for inclusion in the Work
133 | by You to the Licensor shall be under the terms and conditions of
134 | this License, without any additional terms or conditions.
135 | Notwithstanding the above, nothing herein shall supersede or modify
136 | the terms of any separate license agreement you may have executed
137 | with Licensor regarding such Contributions.
138 |
139 | 6. Trademarks. This License does not grant permission to use the trade
140 | names, trademarks, service marks, or product names of the Licensor,
141 | except as required for reasonable and customary use in describing the
142 | origin of the Work and reproducing the content of the NOTICE file.
143 |
144 | 7. Disclaimer of Warranty. Unless required by applicable law or
145 | agreed to in writing, Licensor provides the Work (and each
146 | Contributor provides its Contributions) on an "AS IS" BASIS,
147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
148 | implied, including, without limitation, any warranties or conditions
149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
150 | PARTICULAR PURPOSE. You are solely responsible for determining the
151 | appropriateness of using or redistributing the Work and assume any
152 | risks associated with Your exercise of permissions under this License.
153 |
154 | 8. Limitation of Liability. In no event and under no legal theory,
155 | whether in tort (including negligence), contract, or otherwise,
156 | unless required by applicable law (such as deliberate and grossly
157 | negligent acts) or agreed to in writing, shall any Contributor be
158 | liable to You for damages, including any direct, indirect, special,
159 | incidental, or consequential damages of any character arising as a
160 | result of this License or out of the use or inability to use the
161 | Work (including but not limited to damages for loss of goodwill,
162 | work stoppage, computer failure or malfunction, or any and all
163 | other commercial damages or losses), even if such Contributor
164 | has been advised of the possibility of such damages.
165 |
166 | 9. Accepting Warranty or Additional Liability. While redistributing
167 | the Work or Derivative Works thereof, You may choose to offer,
168 | and charge a fee for, acceptance of support, warranty, indemnity,
169 | or other liability obligations and/or rights consistent with this
170 | License. However, in accepting such obligations, You may act only
171 | on Your own behalf and on Your sole responsibility, not on behalf
172 | of any other Contributor, and only if You agree to indemnify,
173 | defend, and hold each Contributor harmless for any liability
174 | incurred by, or claims asserted against, such Contributor by reason
175 | of your accepting any such warranty or additional liability.
176 |
--------------------------------------------------------------------------------
/NOTICE:
--------------------------------------------------------------------------------
1 | Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # AnoLLM: Large Language Models for Tabular Anomaly Detection (ICLR 2025)
2 |
3 |
4 |
5 | 
6 |
7 |
8 | 
9 |
10 |
11 |
12 | This repository contains the implementation of the paper:
13 | > **AnoLLM: Large Language Models for Tabular Anomaly Detection**
14 | > International Conference on Learning Representations (ICLR 2025)
15 | > Che-Ping Tsai, Ganyu Teng, Phil Wallis, Wei Ding.
16 |
17 | ## Introduction
18 |
19 |
20 |
21 |
22 |
23 |
', '') for _ in text]
446 |
447 | print(f'number of normal samples: {sum(label==0)}, number of anomalies: {sum(label==1)}')
448 |
449 | return text, label
450 | target_label = int(dataset_name.split('-')[1])
451 | text, label = data_generator(subsample=10000, target_label=target_label)
452 | y = label
453 | X = pd.DataFrame(data = text, columns = ['text'])
454 |
455 | elif dataset_name in DATA_MAP.keys():
456 | # datasets from ADBench
457 | dataset_root = Path(adbench.__file__).parent.absolute() / "datasets/Classical"
458 | n = DATA_MAP[dataset_name]
459 | for npz_file in os.listdir(dataset_root):
460 | if npz_file.startswith(str(n) + '_'):
461 | print(dataset_name, npz_file)
462 | data = np.load(dataset_root / npz_file, allow_pickle=False)
463 | break
464 | else:
465 | ValueError('{} is not found.'.format(dataset_name))
466 | X_np, y = data['X'], data['y']
467 | X = convert_np_to_df(X_np)
468 | else:
469 | raise ValueError('Invalid dataset name {}'.format(dataset_name))
470 |
471 | assert len(X) == len(y)
472 |
473 | with open(pkl_file, 'wb') as f:
474 | pickle.dump((X,y), f)
475 |
476 | return X, y
477 |
478 | def load_adbench_data(dataset):
479 | dataset_root = Path(adbench.__file__).parent.absolute() / "datasets/Classical"
480 | if not os.path.exists(dataset_root):
481 | from adbench.myutils import Utils
482 | Utils().download_datasets(repo='jihulab')
483 |
484 | if dataset == 'cardio':
485 | return np.load(dataset_root / '6_cardio.npz', allow_pickle=False)
486 |
487 | for npz_file in os.listdir(dataset_root):
488 | if dataset in npz_file.lower():
489 | return np.load(dataset_root / npz_file, allow_pickle=False)
490 | else:
491 | ValueError('{} is not found.'.format(dataset))
492 |
493 | def split_data(
494 | X: pd.DataFrame,
495 | dataset_name: str,
496 | n_splits: int,
497 | data_dir: str,
498 | train_ratio: Optional[float] = 0.5,
499 | y: Optional[np.ndarray] = None, # should be provided in semi-supervised settinig
500 | seed: Optional[int] = 42,
501 | setting: Optional[str] = 'semi_supervised'
502 | ) -> tuple: # list of train indices and test indices
503 | np.random.seed(seed)
504 | #save path
505 | split_dir = Path(data_dir) / dataset_name / setting / 'split{}'.format(n_splits)
506 | os.makedirs(split_dir, exist_ok = True)
507 |
508 | train_indices, test_indices = [], []
509 | for i in range(n_splits):
510 | pkl_file = split_dir / 'index{}.pkl'.format(i)
511 | if os.path.exists(pkl_file):
512 | with open(pkl_file, 'rb') as f:
513 | train_index, test_index = pickle.load(f)
514 | else:
515 | if setting == 'unsupervised':
516 | normal_data_indices = np.where(y==0)[0]
517 | anormal_data_indices = np.where(y==1)[0]
518 | normal_index = np.random.permutation(normal_data_indices)
519 | anormal_index = np.random.permutation(anormal_data_indices)
520 |
521 | train_index = np.concatenate([normal_index[:int(train_ratio * len(normal_index))], anormal_index[:int(train_ratio * len(anormal_index))]])
522 | test_index = np.concatenate([normal_index[int(train_ratio * len(normal_index)):], anormal_index[int(train_ratio * len(anormal_index)):]])
523 | elif setting == 'semi_supervised':
524 | normal_data_indices = np.where(y==0)[0]
525 | anormal_data_indices = np.where(y==1)[0]
526 | data_length = len(normal_data_indices)
527 | index = np.random.permutation(normal_data_indices)
528 |
529 | train_index = index[:int(train_ratio * data_length)]
530 | test_index = index[int(train_ratio * data_length):]
531 | test_index = np.concatenate([test_index, anormal_data_indices])
532 | else:
533 | raise ValueError('Invalid setting. Choose either unsupervised or semi_supervised')
534 | train_index = np.random.permutation(train_index)
535 | test_index = np.random.permutation(test_index)
536 | with open(pkl_file, 'wb') as f:
537 | pickle.dump((train_index, test_index), f)
538 | train_indices.append(train_index)
539 | test_indices.append(test_index)
540 | return train_indices, test_indices
541 |
542 | def convert_np_to_df(X_np):
543 | n_train, n_cols = X_np.shape
544 | # Add missing column names
545 | L = list(string.ascii_uppercase) + [letter1+letter2 for letter1 in string.ascii_uppercase for letter2 in string.ascii_uppercase]
546 | columns = [ L[i] for i in range(n_cols) ]
547 | df = pd.DataFrame(data = X_np, columns = columns)
548 | return df
549 |
550 | def load_data(args):
551 | dataset_dir = Path(args.data_dir) / args.dataset
552 | X, y = load_dataset(args.dataset, args.data_dir)
553 |
554 | if 'binning' in args and args.binning != 'none':
555 | X = normalize(X, args.binning, args.n_buckets)
556 | if 'remove_feature_name' in args and args.remove_feature_name:
557 | print("Removing column names and category names.")
558 | L = list(ascii_uppercase) + [letter1+letter2 for letter1 in ascii_uppercase for letter2 in ascii_uppercase]
559 | X.columns = [ L[i] for i in range(len(X.columns))]
560 |
561 | categorical_data = X.select_dtypes(include = ['object'])
562 | categorical_columns = categorical_data.columns.tolist()
563 | le = LabelEncoder()
564 | for i in categorical_data.columns:
565 | categorical_data[i] = le.fit_transform(categorical_data[i])
566 |
567 | X_prime = X.drop(categorical_columns, axis = 1)
568 | X = pd.concat([X_prime, categorical_data], axis = 1)
569 |
570 | if 'train_ratio' not in args:
571 | args.train_ratio = 0.5
572 | if 'seed' not in args:
573 | args.seed = 42
574 | train_indices, test_indices = split_data(X, args.dataset, args.n_splits, args.data_dir,
575 | args.train_ratio, y = y, seed = args.seed, setting = args.setting )
576 | train_index, test_index = train_indices[args.split_idx], test_indices[args.split_idx]
577 | X_train, X_test = X.loc[train_index], X.loc[test_index]
578 | y_train, y_test = y[train_index], y[test_index]
579 |
580 | return X_train, X_test, y_train, y_test
581 |
582 | def normalize(X, method, n_buckets):
583 | # method: ['quantile', 'equal_width', 'language', 'none', 'standard']
584 | # n_buckets: 0-100
585 | X = copy.deepcopy(X)
586 | def ordinal(n):
587 | if np.isnan(n):
588 | return 'NaN'
589 | n = int(n)
590 | if 10 <= n % 100 <= 20:
591 | suffix = 'th'
592 | else:
593 | suffix = {1: 'st', 2: 'nd', 3: 'rd'}.get(n % 10, 'th')
594 | return 'the ' + str(n) + suffix + ' percentile'
595 |
596 | word_list = ['Minimal', 'Slight', 'Moderate', 'Noticeable', 'Considerable', 'Significant', 'Substantial', 'Major', 'Extensive', 'Maximum']
597 | def get_word(n):
598 | n = int(n)
599 | if n == 10:
600 | return word_list[-1]
601 | return word_list[n]
602 |
603 | if method == 'quantile':
604 | for column in X.columns:
605 | if X[column].dtype in ['float64', 'int64', 'uint8', 'int16'] and X[column].nunique() > 1:
606 | ranks = X[column].rank(method='min')
607 | X[column] = ranks / len(X[column]) * 100
608 | X[column] = X[column].apply(ordinal)
609 |
610 | elif method == 'equal_width':
611 | for column in X.columns:
612 | if X[column].dtype in ['float64', 'int64', 'uint8', 'int16']:
613 | if X[column].nunique() > 1:
614 | X[column] = X[column].astype('float64')
615 | X[column] = (X[column] - X[column].min()) / (X[column].max() - X[column].min()) * n_buckets
616 |
617 | if 10 % n_buckets == 0:
618 | X[column] = X[column].round(0) / 10
619 | X[column] = X[column].round(1)
620 | else:
621 | X[column] = X[column].round(0) / 100
622 | X[column] = X[column].round(2)
623 | elif method == 'standard':
624 | for column in X.columns:
625 | if X[column].dtype in ['float64', 'int64', 'uint8', 'int16']:
626 | scaler = StandardScaler()
627 | scaler.fit(X[column].values.reshape(-1,1))
628 | X[column] = scaler.transform(X[column].values.reshape(-1,1))
629 | X[column] = X[column].round(1)
630 |
631 | elif method == 'language':
632 | for column in X.columns:
633 | if X[column].dtype in ['float64', 'int64', 'uint8', 'int16'] and X[column].nunique() > 1:
634 | X[column] = X[column].astype('float64')
635 | X[column] = (X[column] - X[column].min()) / (X[column].max() - X[column].min()) * 10
636 | X[column] = X[column].apply(get_word)
637 | else:
638 | raise ValueError('Invalid method. Choose either percentile, language or decimal')
639 | return X
640 |
641 | def get_text_columns(dataset_name):
642 | text_columns = []
643 | if dataset_name == 'fakejob':
644 | text_columns = ['title', 'company profile', 'description', 'requirements', 'benefits']
645 | elif 'fakenews' == dataset_name:
646 | text_columns = ['title', 'text']
647 | elif '20news' in dataset_name:
648 | text_columns = ['text']
649 | return text_columns
650 |
651 | def get_max_length_dict(dataset_name):
652 | max_length_dict = {}
653 | if dataset_name == 'fakejob':
654 | max_length_dict['title'] = 20
655 | text_columns = ['company profile', 'description', 'requirements', 'benefits']
656 | for col in text_columns:
657 | max_length_dict[col] = 700
658 | elif 'fakenews' == dataset_name:
659 | max_length_dict['title'] = 30
660 | max_length_dict['text'] = 500
661 | elif '20news' in dataset_name:
662 | max_length_dict['text'] = 1000
663 | return max_length_dict
664 |
665 | def df_to_numpy(
666 | X: pd.DataFrame,
667 | dataset_name: Optional[str] = None,
668 | method: Optional[str] = 'ordinal',
669 | normalize_numbers: Optional[bool] = False,
670 | verbose: Optional[bool] = False,
671 | textual_encoding: Optional[str] = 'word2vec', # bag_of_words, tfidf, word2vec, or none
672 | textual_columns: Optional[list] = None
673 | ) -> np.ndarray:
674 | if dataset_name == 'ecoli':
675 | X_np = X.drop(X.columns[0], axis=1).to_numpy()
676 | return X_np
677 |
678 | numeric_data = X.select_dtypes(include = ['float64', 'int64', 'uint8', 'int16', 'float32'])
679 | numeric_columns = numeric_data.columns.tolist()
680 | categorical_data = X.select_dtypes(include = ['object', 'category'])
681 | categorical_columns = categorical_data.columns.tolist()
682 |
683 | if verbose:
684 | print("Number of categorical data", len(categorical_columns))
685 | print("Categorical columns:", categorical_columns)
686 |
687 | # fill na
688 | if len(numeric_columns) > 0:
689 | for numeric_col in numeric_columns:
690 | X[numeric_col] = X[numeric_col].fillna(X[numeric_col].mean())
691 |
692 | if normalize_numbers:
693 | # normalize it to have zero mean and unit variance
694 | scaler = StandardScaler()
695 | X[numeric_columns] = scaler.fit_transform(X[numeric_columns])
696 |
697 | # Handle textual data
698 | if textual_encoding == 'none' and len(textual_columns) > 0:
699 | for col in textual_columns:
700 | categorical_columns.remove(col)
701 | X = X.drop(columns = textual_columns)
702 | textual_columns = []
703 |
704 | if len(textual_columns) > 0:
705 | if textual_encoding == 'word2vec':
706 | model = api.load('word2vec-google-news-300')
707 | tmp = X[textual_columns].agg(' '.join, axis=1)
708 | X_vecs = []
709 | for i in range(len(X)):
710 | words = []
711 | for word in tmp[i].split():
712 | if word in model.key_to_index:
713 | words.append(word)
714 | # Compute the average word embedding
715 | if words: # Ensure there are valid words left
716 | word_vectors = [model[word] for word in words]
717 | X_vec = np.mean(word_vectors, axis=0)
718 | else:
719 | X_vec = np.zeros(model.vector_size) # Handle the case where no words are in the vocabulary
720 | X_vecs.append(X_vec)
721 | X_vecs = np.array(X_vecs)
722 | for col in textual_columns:
723 | categorical_columns.remove(col)
724 |
725 | elif textual_encoding == 'bag_of_words':
726 | corpus = []
727 | for col in textual_columns:
728 | for i in range(len(X)):
729 | corpus.append(X[col][i])
730 | vectorization = CountVectorizer(max_features = 300)
731 | vectorization.fit(corpus)
732 | tmp = X[textual_columns].agg(' '.join, axis=1)
733 | X_vecs = vectorization.transform(tmp).todense()
734 |
735 | for col in textual_columns:
736 | categorical_columns.remove(col)
737 |
738 | elif textual_encoding == 'tfidf':
739 | corpus = []
740 | for col in textual_columns:
741 | for i in range(len(X)):
742 | corpus.append(X[col][i])
743 | vectorization = TfidfVectorizer(max_features = 300)
744 | vectorization.fit(corpus)
745 | tmp = X[textual_columns].agg(' '.join, axis=1)
746 | X_vecs = vectorization.transform(tmp).todense()
747 |
748 | for col in textual_columns:
749 | categorical_columns.remove(col)
750 |
751 | else:
752 | raise ValueError('Invalid textual encoding. Choose either bag_of_words, tf-idf or word2vec')
753 | X = X.drop(columns = textual_columns)
754 | X = pd.concat([X, pd.DataFrame(X_vecs)], axis = 1)
755 |
756 |
757 | if len(categorical_columns) > 0:
758 | # categorical features:
759 | # group categories with low frequency into a single category
760 | encoder = RareLabelEncoder(
761 | tol=0.01, # Minimum frequency to be considered as a separate class
762 | max_n_categories=None, # Maximum number of categories to keep
763 | replace_with='Rare', # Value to replace rare categories with
764 | variables=categorical_columns , # Columns to encode
765 | missing_values='ignore',
766 | )
767 | X = encoder.fit_transform(X)
768 |
769 | # Remove columns that contain identical values
770 | X = X.loc[:, (X != X.iloc[0]).any()]
771 |
772 | # remove categories that have only one value
773 | for column in categorical_columns:
774 | if X[column].nunique() == 1:
775 | X.drop(column, inplace = True, axis = 1)
776 |
777 | if method == 'ordinal':
778 | le = LabelEncoder()
779 | for i in categorical_data.columns:
780 | categorical_data[i] = le.fit_transform(categorical_data[i])
781 | elif method == 'one_hot':
782 | enc = OneHotEncoder(handle_unknown='ignore', sparse_output=False, drop='first')
783 | one_hot_encoded = enc.fit_transform(X[categorical_columns])
784 | categorical_data = pd.DataFrame(one_hot_encoded, columns=enc.get_feature_names_out(categorical_columns))
785 | else:
786 | raise ValueError('Invalid method. Choose either ordinal or one_hot')
787 | X_prime = X.drop(categorical_columns, axis = 1)
788 | X = pd.concat([X_prime, categorical_data], axis = 1)
789 | # remove columns that contain identical values
790 | print(X.shape)
791 | X = X.loc[:, (X != X.iloc[0]).any()]
792 | X_np = X.to_numpy()
793 | return X_np
794 |
795 | def print_dataset_information(dataset, data_dir):
796 | print("-"*100)
797 | print("Dataset: {}".format(dataset))
798 | X, y = load_dataset(dataset, data_dir)
799 | #print(X['company profile'][:3])
800 | print(X.columns)
801 | train_indices, test_indices = split_data(X, dataset, 5, data_dir,
802 | 0.5, y = y, seed = 42, setting = 'semi_supervised' )
803 | print("Dtypes of columns:", X.dtypes)
804 | X_np = df_to_numpy(X, dataset_name = dataset, method = 'one_hot', verbose = True, textual_encoding='word2vec')
805 | print("Number of training samples:", len(train_indices[0]))
806 | print("Number of testing samples:", len(test_indices[0]))
807 | print("Number of anomalies: {:f} ({:.2f}%)".format(np.sum(y), np.sum(y)/len(y) * 100))
808 | print("Number of features:", len(X.columns))
809 | print("Number of feature dimensions", X_np.shape[1])
810 |
811 | def filter_anomalies(X_test, y_test):
812 | X_test = X_test[y_test == 0]
813 | y_test = y_test[y_test == 0]
814 | return X_test, y_test
815 |
816 |
817 | if __name__ == '__main__':
818 | #print_dataset_information('fakejob', 'data')
819 | print_dataset_information('20news-6', 'data')
820 | exit()
--------------------------------------------------------------------------------
/src/get_avg_results.py:
--------------------------------------------------------------------------------
1 | import argparse
2 |
3 | import numpy as np
4 | import pandas as pd
5 |
6 | from data_utils import DATA_MAP, MIXED, ODDS
7 | from get_results import get_metrics, aggregate_results, filter_results
8 |
9 | def get_args():
10 | parser = argparse.ArgumentParser()
11 | parser.add_argument("--dataset", type = str, default='all', choices = ['all', 'mixed', 'odds']) # subset: 10 datasets that contain feature names
12 | parser.add_argument("--setting", type = str, default='semi_supervised', choices = ['semi_supervised', 'unsupervised'])
13 | parser.add_argument("--exp_dir", type = str, default=None)
14 | parser.add_argument("--metric", type = str, choices =["AUC-ROC", "F1", "AUC-PR"], default='AUC-ROC')
15 | #dataset hyperparameters
16 | parser.add_argument("--data_dir", type = str, default='data')
17 | parser.add_argument("--n_splits", type = int, default=5)
18 | parser.add_argument("--only_normalized", action='store_true', default=False)
19 | parser.add_argument("--only_ordinal", action='store_true', default=False)
20 | args = parser.parse_args()
21 |
22 | return args
23 |
24 | def main():
25 | args = get_args()
26 | roc_scores = {}
27 | ranking_dict = {}
28 | std_scores = {}
29 | if args.dataset == 'all':
30 | DATASETS = [k for k in DATA_MAP.keys()]
31 | elif args.dataset == 'odds':
32 | DATASETS = ODDS
33 | elif args.dataset == 'mixed':
34 | DATASETS = MIXED
35 | # sorted datasets by alphabetic order
36 | DATASETS = sorted(DATASETS)
37 | all_rocs = {}
38 | for dataset_idx, dataset in enumerate(DATASETS):
39 | try:
40 | print("*"*100)
41 | print(dataset)
42 | args.split_idx = None
43 | args.dataset = dataset
44 | L = []
45 | for i in range(args.n_splits):
46 | args.split_idx = i
47 | args.exp_dir = None
48 | results = get_metrics(args, only_normalized=args.only_normalized, only_ordinal=args.only_ordinal)
49 | L.append(results)
50 | metrics, rankings = aggregate_results(L)
51 | except:
52 | print("Error in dataset: ", dataset)
53 | continue
54 |
55 | metrics = filter_results(metrics)
56 | for k in metrics.keys():
57 |
58 | if k not in all_rocs:
59 | all_rocs[k] = np.zeros((len(DATASETS), args.n_splits))
60 |
61 | for i in range(args.n_splits):
62 | all_rocs[k][dataset_idx] = metrics[k][args.metric]
63 |
64 | roc_scores[dataset] = { k: np.mean(metrics[k][args.metric]) for k in metrics.keys()}
65 | ranking_dict[dataset] = {k: int(rankings[args.metric][idx]) for idx, k in enumerate(metrics.keys())}
66 | std_scores[dataset] = { k: np.std(metrics[k][args.metric]) for k in metrics.keys()}
67 |
68 | df = pd.DataFrame(roc_scores).T
69 | avg_row = df.mean(axis=0)
70 | df.loc['avg'] = avg_row
71 | df = df.round(3)
72 | print(df)
73 | df.to_csv('exp/{}_avg_{}.csv'.format(args.setting, args.metric))
74 |
75 | '''
76 | ranking_df = pd.DataFrame(ranking_dict).T
77 | avg_row = ranking_df.mean(axis=0)
78 | ranking_df.loc['avg'] = avg_row
79 | ranking_df = ranking_df.round(3)
80 | print(ranking_df)
81 | ranking_df.to_csv('exp/{}_avg_ranking.csv'.format(args.setting))
82 | '''
83 | # std
84 | std_df = pd.DataFrame(std_scores).T
85 | avg_std = []
86 | for c in std_df.columns:
87 | if c not in all_rocs:
88 | avg_std.append(0)
89 | continue
90 | std = np.std( np.mean(all_rocs[c], axis=0))
91 | avg_std.append(std)
92 | std_df.loc['avg'] = avg_std
93 | std_df = std_df.round(3)
94 | print(std_df)
95 | std_df.to_csv('exp/{}_std_{}.csv'.format(args.setting,args.metric))
96 |
97 | if __name__ == '__main__':
98 | main()
99 |
--------------------------------------------------------------------------------
/src/get_results.py:
--------------------------------------------------------------------------------
1 | import os
2 | from pathlib import Path
3 | import argparse
4 |
5 | from sklearn import metrics
6 | import numpy as np
7 | import pandas as pd
8 | from data_utils import load_data, DATA_MAP
9 |
10 |
11 | def get_args():
12 | parser = argparse.ArgumentParser()
13 | parser.add_argument("--dataset", type = str, default='wine', choices = [d.lower() for d in DATA_MAP.keys()],
14 | help="Name of datasets in the ODDS benchmark")
15 | parser.add_argument("--exp_dir", type = str, default=None)
16 | parser.add_argument("--setting", type = str, default='semi_supervised', choices = ['semi_supervised', 'unsupervised'])
17 |
18 | #dataset hyperparameters
19 | parser.add_argument("--data_dir", type = str, default='data')
20 | parser.add_argument("--n_splits", type = int, default=5)
21 | parser.add_argument("--split_idx", type = int, default=None) # 0 to n_split-1
22 |
23 | args = parser.parse_args()
24 |
25 | return args
26 |
27 | def tabular_metrics(y_true, y_score):
28 | """
29 | Calculates evaluation metrics for tabular anomaly detection.
30 | Adapted from https://github.com/xuhongzuo/DeepOD/blob/main/deepod/metrics/_anomaly_detection.py
31 | Args:
32 |
33 | y_true (np.array, required):
34 | Data label, 0 indicates normal timestamp, and 1 is anomaly.
35 |
36 | y_score (np.array, required):
37 | Predicted anomaly scores, higher score indicates higher likelihoods to be anomaly.
38 |
39 | Returns:
40 | tuple: A tuple containing:
41 |
42 | - auc_roc (float):
43 | The score of area under the ROC curve.
44 |
45 | - auc_pr (float):
46 | The score of area under the precision-recall curve.
47 |
48 | - f1 (float):
49 | The score of F1-score.
50 |
51 | - precision (float):
52 | The score of precision.
53 |
54 | - recall (float):
55 | The score of recall.
56 |
57 | """
58 | # F1@k, using real percentage to calculate F1-score
59 | n_test = len(y_true)
60 | new_index = np.random.permutation(n_test) # shuffle y to prevent bias of ordering (argpartition may discard entries with same value)
61 | y_true = y_true[new_index]
62 | y_score = y_score[new_index]
63 |
64 | #ratio = 100.0 * len(np.where(y_true == 0)[0]) / len(y_true)
65 | #thresh = np.percentile(y_score, ratio)
66 | #y_pred = (y_score >= thresh).astype(int)
67 |
68 | top_k = len(np.where(y_true == 1)[0])
69 | indices = np.argpartition(y_score, -top_k)[-top_k:]
70 | y_pred = np.zeros_like(y_true)
71 | y_pred[indices] = 1
72 |
73 | y_true = y_true.astype(int)
74 | p, r, f1, support = metrics.precision_recall_fscore_support(y_true, y_pred, average='binary')
75 |
76 | return metrics.roc_auc_score(y_true, y_score), metrics.average_precision_score(y_true, y_score), f1, p, r
77 |
78 | def get_metrics(args, only_raw = False, only_normalized = False, only_ordinal = False):
79 | X_train, X_test, y_train, y_test = load_data(args)
80 | if isinstance(y_test, pd.Series):
81 | y_test = np.array(y_test)
82 | #y_test = y_test.to_numpy()
83 | if args.exp_dir is None:
84 | args.exp_dir = Path('exp') / args.dataset / args.setting / "split{}".format(args.n_splits) / "split{}".format(args.split_idx)
85 | score_dir = args.exp_dir / 'scores'
86 | if not os.path.exists(score_dir):
87 | raise ValueError("Score directory {} does not exist".format(score_dir))
88 |
89 |
90 | method_dict = {}
91 | for score_npy in os.listdir(score_dir):
92 | if '.npy' in score_npy:
93 | if score_npy.startswith('raw'):
94 | continue
95 | if is_baseline(score_npy) and only_normalized:
96 | if 'normalized' not in score_npy:
97 | continue
98 |
99 | elif is_baseline(score_npy) and only_ordinal:
100 | if 'ordinal' not in score_npy:
101 | continue
102 |
103 |
104 | method = '.'.join(score_npy.split('.')[:-1])
105 | if method == 'rdp':
106 | continue
107 | scores = np.load(score_dir / score_npy)
108 | if np.isnan(scores).any():
109 | print("NaNs in scores for {}".format(method))
110 | method_dict[method] = [0, 0, 0, 0, 0]
111 | elif np.isinf(scores).any():
112 | print("Infs in scores for {}".format(method))
113 | method_dict[method] = [0, 0, 0, 0, 0]
114 | else:
115 | auc_roc, auc_pr, f1, p, r = tabular_metrics(y_test, scores)
116 | method_dict[method] = [auc_roc, auc_pr, f1, p, r]
117 | # get ranking info for all methods
118 | rankings = []
119 | method = list(method_dict.keys())[0]
120 | for i in range(len(method_dict[method])):
121 | scores = [-method_dict[k][i] for k in method_dict.keys()]
122 | ranking = np.argsort(scores).argsort() + 1
123 | rankings.append(ranking)
124 |
125 | print("-"*100)
126 | for idx, (k, v) in enumerate(method_dict.items()):
127 | #print("{:30s}: AUC-ROC: {:.4f}, AUC-PR: {:.4f}, F1: {:.4f}".format(k, v[0], v[1], v[2]))
128 | print("{:30s}: AUC-ROC: {:.4f} ({:2d}), AUC-PR: {:.4f} ({:2d}), F1: {:.4f} ({:2d}), P: {:.4f} ({:2d}), R: {:.4f} ({:2d})".format(k,
129 | v[0], rankings[0][idx],
130 | v[1], rankings[1][idx],
131 | v[2], rankings[2][idx],
132 | v[3], rankings[3][idx],
133 | v[4], rankings[4][idx],
134 | ))
135 |
136 | return method_dict
137 | def is_baseline(s):
138 | if 'anollm' in s:
139 | return False
140 | return True
141 |
142 | def filter_results(d:dict):
143 | d2 = {}
144 | for k in d.keys():
145 | new_key = k
146 | if is_baseline(k):
147 | #baselines
148 | d2[k] = d[k]
149 | else:
150 | if '_lora' in k:
151 | temp = k.replace('_lora', '')
152 | if temp in d:
153 | continue
154 | else:
155 | new_key = new_key.replace('_lora', '')
156 | d2[new_key] = d[k]
157 | else:
158 | d2[new_key] = d[k]
159 | return d2
160 |
161 | def aggregate_results(m_dicts):
162 | aggregate_results = {k: {'AUC-ROC':[], 'AUC-PR': [], 'F1': [], 'P': [], 'R':[]} for k in m_dicts[0].keys()}
163 | for i in range(len(m_dicts)):
164 | all_keys = list(m_dicts[0].keys())
165 | for k in all_keys:
166 | try:
167 | aggregate_results[k]['AUC-ROC'] += [m_dicts[i][k][0]]
168 | aggregate_results[k]['AUC-PR'] += [m_dicts[i][k][1]]
169 | aggregate_results[k]['F1'] += [m_dicts[i][k][2]]
170 | aggregate_results[k]['P'] += [m_dicts[i][k][3]]
171 | aggregate_results[k]['R'] += [m_dicts[i][k][4]]
172 | except:
173 | print("Incomplete results for ", k)
174 | if k in aggregate_results:
175 | del aggregate_results[k]
176 | for i in range(len(m_dicts)):
177 | if k in m_dicts[i]:
178 | del m_dicts[i][k]
179 | continue
180 |
181 | print("-"*100)
182 |
183 | # get ranking info for all methods
184 | rankings = {}
185 | key = list(m_dicts[0].keys())[0]
186 | for metric_name in aggregate_results[key].keys():
187 | scores = [-np.mean(aggregate_results[k][metric_name]) for k in aggregate_results.keys()]
188 | ranking = np.argsort(scores).argsort() + 1
189 | rankings[metric_name] = ranking
190 |
191 | for idx, k in enumerate(aggregate_results.keys()):
192 | print("{:30s}: AUC-ROC: {:.4f} +- {:.4f} ({:2d}), AUC-PR: {:.4f} +- {:.4f} ({:2d}), F1: {:.4f} +- {:.4f} ({:2d}) P: {:.4f} +- {:.4f} ({:2d}) R: {:.4f} +- {:.4f} ({:2d})".format(k,
193 | np.mean(aggregate_results[k]['AUC-ROC']), np.std(aggregate_results[k]['AUC-ROC']), rankings['AUC-ROC'][idx],
194 | np.mean(aggregate_results[k]['AUC-PR']), np.std(aggregate_results[k]['AUC-PR']), rankings['AUC-PR'][idx],
195 | np.mean(aggregate_results[k]['F1']), np.std(aggregate_results[k]['F1']), rankings['F1'][idx],
196 | np.mean(aggregate_results[k]['P']), np.std(aggregate_results[k]['P']), rankings['P'][idx],
197 | np.mean(aggregate_results[k]['R']), np.std(aggregate_results[k]['R']), rankings['R'][idx],
198 | ))
199 | return aggregate_results, rankings
200 |
201 | def main():
202 | args = get_args()
203 | if args.split_idx is None:
204 | L = []
205 | for i in range(args.n_splits):
206 | args.split_idx = i
207 | args.exp_dir = None
208 | results = get_metrics(args)
209 | L.append(results)
210 | aggregate_results(L)
211 | else:
212 | print(args)
213 | scores = get_metrics(args)
214 |
215 |
216 | if __name__ == '__main__':
217 | main()
--------------------------------------------------------------------------------
/train_anollm.py:
--------------------------------------------------------------------------------
1 | import os
2 | from pathlib import Path
3 |
4 | import numpy as np
5 | import argparse
6 | import torch.distributed as dist
7 | import torch
8 | import wandb
9 | import time
10 |
11 | from anollm import AnoLLM
12 | from src.data_utils import load_data, DATA_MAP, get_text_columns, get_max_length_dict
13 |
14 | #run by torchrun --nproc_per_node=8 train_llm.py
15 |
16 | def get_args():
17 | parser = argparse.ArgumentParser()
18 | parser.add_argument("--dataset", type = str, default='wine', choices = [d.lower() for d in DATA_MAP.keys()],
19 | help="Name of datasets in the ODDS benchmark")
20 | parser.add_argument("--exp_dir", type = str, default=None)
21 | parser.add_argument("--setting", type = str, default='semi_supervised', choices = ['semi_supervised', 'unsupervised'], help="semi_supervised:an uncontaminated, unsupervised setting; unsupervised:a contaminated, unsupervised setting")
22 |
23 | # wandb
24 | parser.add_argument("--wandb", action='store_true')
25 | parser.add_argument("--entity", type = str, default = None)
26 | parser.add_argument("--project", type = str, default = 'AnoLLM')
27 |
28 | #dataset hyperparameters
29 | parser.add_argument("--data_dir", type = str, default='data')
30 | parser.add_argument("--n_splits", type = int, default=5)
31 | parser.add_argument("--split_idx", type = int, default=0) # 0 to n_split-1
32 | parser.add_argument("--train_ratio", type = float, default=0.5)
33 | parser.add_argument("--seed", type = int, default=42)
34 |
35 | # preprocessing
36 | parser.add_argument("--binning", type = str, choices=['quantile', 'equal_width', 'language', 'none', 'standard'], default='standard')
37 | parser.add_argument("--n_buckets", type = int, default=10)
38 | parser.add_argument("--remove_feature_name", action = 'store_true')
39 |
40 | #training
41 | parser.add_argument("--model", type = str, choices = ['gpt2', 'distilgpt2', 'smol', 'smol-360', 'smol-1.7b'], default='smol')
42 | parser.add_argument("--batch_size", type = int, default=32) # per gpu, eval_batch_size = 2*batch_size
43 | parser.add_argument("--lr", type = float, default=5e-5)
44 | parser.add_argument("--lora", action='store_true', default=False)
45 | parser.add_argument("--max_steps", type = int, default=2000)
46 | parser.add_argument("--eval_steps", type = int, default = 1000)
47 | parser.add_argument("--random_init", action='store_true', default=False)
48 | parser.add_argument("--no_random_permutation", action='store_true', default=False)
49 |
50 | args = parser.parse_args()
51 | if args.exp_dir is None:
52 | args.exp_dir = Path('exp') / args.dataset / args.setting / "split{}".format(args.n_splits) / "split{}".format(args.split_idx)
53 |
54 | if args.model == 'smol':
55 | args.model = 'HuggingFaceTB/SmolLM-135M'
56 | elif args.model == 'smol-360':
57 | args.model = 'HuggingFaceTB/SmolLM-360M'
58 | elif args.model == 'smol-1.7b':
59 | args.model = 'HuggingFaceTB/SmolLM-1.7B'
60 |
61 | args.save_dir = Path(args.exp_dir) / 'models' # save to save models
62 | os.makedirs(args.save_dir, exist_ok = True)
63 |
64 | return args
65 |
66 | def get_run_name(args):
67 | name = 'anollm'
68 | name += '_lr{}'.format(args.lr)
69 | name += '_{}'.format(args.binning)
70 |
71 | if args.model == 'HuggingFaceTB/SmolLM-135M':
72 | name += '_smolLM'
73 | elif args.model == 'HuggingFaceTB/SmolLM-360M':
74 | name += '_smolLM360'
75 | elif args.model == 'HuggingFaceTB/SmolLM-1.7B':
76 | name += '_smolLM1.7B'
77 | else:
78 | name += '_' + args.model
79 |
80 | if args.random_init:
81 | name += '_random_init'
82 |
83 | if args.no_random_permutation:
84 | name += '_no_random_permutation'
85 |
86 | if args.lora:
87 | name += '_lora'
88 | name += "_test"
89 | return name
90 |
91 |
92 | def main():
93 | # Set CUDA devices for each process
94 | local_rank = int(os.environ["LOCAL_RANK"])
95 | torch.cuda.set_device(local_rank)
96 |
97 | args = get_args()
98 | if dist.get_rank() == 0:
99 | X_train, X_test, y_train, y_test = load_data(args)
100 | dist.barrier()
101 | if dist.get_rank() != 0:
102 | X_train, X_test, y_train, y_test = load_data(args)
103 | dist.barrier()
104 |
105 | run_name = get_run_name(args)
106 | efficient_finetuning = 'lora' if args.lora else ''
107 | model_path = args.save_dir / '{}.pt'.format(run_name)
108 | dataset_tmp_path = args.save_dir / (run_name + '_data')
109 |
110 | os.makedirs(dataset_tmp_path, exist_ok= True)
111 | print("Model path:", model_path)
112 | #if False:
113 | if os.path.exists(model_path):
114 | print("Model exists, skip training")
115 | return
116 |
117 | max_length_dict = get_max_length_dict(args.dataset)
118 | text_columns = get_text_columns(args.dataset)
119 | def get_model():
120 | model = AnoLLM(args.model,
121 | batch_size=args.batch_size,
122 | max_steps = args.max_steps,
123 | efficient_finetuning = efficient_finetuning,
124 | max_length_dict=max_length_dict,
125 | textual_columns = text_columns,
126 | random_init=args.random_init,
127 | no_random_permutation=args.no_random_permutation,
128 | bf16=True,
129 | adam_beta2=0.99,
130 | adam_epsilon=1e-7,
131 | learning_rate=args.lr,
132 | )
133 | return model
134 | # Initialize the LLM
135 | if dist.get_rank() == 0:
136 | anollm = get_model()
137 | dist.barrier()
138 | if dist.get_rank() != 0:
139 | anollm = get_model()
140 | dist.barrier()
141 | # Move the model to the appropriate GPU
142 | anollm.model.to(local_rank)
143 |
144 | # Wrap the model for distributed training
145 | anollm.model = torch.nn.parallel.DistributedDataParallel(
146 | anollm.model, device_ids=[local_rank], output_device=local_rank
147 | )
148 | if args.wandb and dist.get_rank() == 0:
149 | run = wandb.init(
150 | entity=args.entity,
151 | project=args.project,
152 | name = "{}_splits{}_{}_{}".format(args.dataset, args.split_idx, args.n_splits, run_name),
153 | )
154 | if len(X_test) > 3000:
155 | np.random.seed(args.seed)
156 | X_test.reset_index(drop = True, inplace = True)
157 | indices = np.random.choice(len(X_test), 3000, replace = False)
158 | X_test = X_test.loc[indices].reset_index(drop = True)
159 | y_test = y_test[indices]
160 | if not args.wandb:
161 | X_test, y_test = None, None
162 |
163 | # Train the model
164 | start_time = time.time()
165 | trainer = anollm.fit(X_train, X_train.columns.to_list(),
166 | use_wandb = args.wandb,
167 | data_val=X_test,
168 | label_val = y_test,
169 | eval_steps = args.eval_steps,
170 | processed_data_dir = dataset_tmp_path,
171 | )
172 | end_time = time.time()
173 |
174 | # Save the model only from rank 0 process
175 | if dist.get_rank() == 0:
176 |
177 | print("Training time:", end_time - start_time)
178 | run_time_dir = args.exp_dir / "run_time" / "train"
179 | os.makedirs(run_time_dir, exist_ok = True)
180 | run_time_path = run_time_dir / "{}.txt".format(run_name)
181 | with open(run_time_path, 'w') as f:
182 | f.write(str(end_time - start_time))
183 |
184 | print("Save model to ", model_path)
185 | anollm.save_state_dict(model_path)
186 |
187 |
188 | dist.destroy_process_group()
189 |
190 | if __name__ == "__main__":
191 | # Initialize the distributed process group
192 | dist.init_process_group(backend="nccl")
193 | main()
--------------------------------------------------------------------------------