├── LICENSE ├── README.md ├── default-mapper.yaml ├── default.yaml ├── example_phrases ├── Decensoring-autoweights.yaml ├── Decensoring-weighted.yaml └── GPT-Roleplay.yaml ├── images ├── cyclic.png ├── gradient.png └── slice.png ├── merge-monster.py ├── modules ├── composition.py ├── logo.ascii ├── mapping.py ├── merging.py ├── models.py ├── probability.py └── utils.py └── monster-mapper.py /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # MergeMonster 2 | An unsupervised merging algorithm for Transformers-based language models, using a list of phrases (both good and bad) and a fully automated strategy that strives to decrease (or increase) the probability of these phrases occurring in the final merge. 3 | 4 | **Refer to the default.yaml example configuration for an explanation of all potential configuration options.** 5 | 6 | ## NEW: MonsterMapper 7 | MonsterMapper is a companion tool to MergeMonster that reuses its configuration files to check the first model's probabilities for the defined contexts. At its core it basically just lists the most probable auto-completions for the contexts you entered in your phrase dictionaries so that you can make the necessary changes to increase the algorithm's efficiency. 8 | 9 | Perhaps the phrase you originally defined isn't as common as you thought, but there's a much more common auto-completion that matches your intended goal. MonsterMapper will help you discover this. 10 | 11 | A seperate config (default-mapper.yaml) has been made available as this tool has its own set of exclusive options to configure as needed. 12 | 13 | ## How It Works 14 | 15 | 1. The algorithm loads a base model of your own choosing (model 1), along with a directory (or repo list) containing multiple models of the same architecture and size. 16 | 2. Each model from the directory is loaded one by one and merged with model 1 on a layer-by-layer basis. 17 | 3. For each layer merge, the algorithm evaluates whether the merge improves the base model, iterating along a customizable list of merging ratios. (See YAML) 18 | - If the merge is beneficial (it lowers the cumulative probability), it is permanently applied to model 1. 19 | - If not, model 1 retains its current structure, and the algorithm proceeds to the next layer. 20 | 4. Upon completing the merge process with one model, the algorithm proceeds to the next model in the list, repeating the cycle. 21 | 5. After all models have been processed, the algorithm saves the final merged model and generates a complete log of the entire process. 22 | 23 | At its very core Merge Monster is nothing more but a relentless number chaser - It tries to decrease the probability of unwanted completions. Wanted completions also subtract from that same number (the monster only cares about lowering the total number), which is why the number displayed during processing might go negative. 24 | 25 | ## Progress bar layout 26 | ``` 27 | [[1.0, 'Mistral-7B-v0.1'], [0.5, 'Nous-Capybara-7B-V1.9'], [0.25, 'SynthIA-7B-v1.3'], [0.1, 'zephyr-7b-beta']] 28 | 09:19:57 - Layer 3/32 - CHANGED - 0.38319 > 0.38261 - 0.2% 29 | 30 | [List of merges applied to this layer, with the first being model 1] 31 | Current time - Layer progress - CHANGED/RETAINED - Old total probability > New total probability - Global change in percentage 32 | ``` 33 | ## Merge Methods 34 | 35 | New merging methods have been added (slice/cyclic) to help target specific parts of tensors. These are highly experimental but verified to be fully functional. 36 | 37 | **Multiple merging methods**: It is possible to combine multiple merging methods by supplying a list to the `merge_method` parameter, such as `["lerp","cyclic"]`. The loop will look like `LAYER > LERP > CYCLIC > NEXT LAYER`. 38 | 39 | **The following methods are currently available:** 40 | 41 | ### "lerp" 42 | 43 | Default method. Linear Interpolation, your basic merging method. 44 | 45 | ### "slerp" 46 | 47 | - Spherical Linear Interpolation, which better aligns the weights inside the two tensors. 48 | - Full credit to [Charles Goddard's mergekit](https://github.com/cg123/mergekit) for the Slerp function. 49 | 50 | ### "slice" 51 | 52 | - Highly experimental. A shift from one model to another, with a smooth 10% transition in-between. 53 | - Ratio defines the starting point of the transition. 54 | 55 | ![Slice](images/slice.png?raw=true "Slice") 56 | 57 | ### "cyclic" 58 | 59 | - Highly experimental. A shift from one model to the other, then back to the original model. 60 | - Especially useful for when you wish to target specific parts of model 1's tensors as model 2 only has a 15% contribution in the resulting output tensor. 61 | - Merge ratios are ignored and a predefined scale is used that covers the entire spectrum during the optimization process. 62 | 63 | ![Cyclic](images/cyclic.png?raw=true "Cyclic") 64 | 65 | ### "gradient" 66 | 67 | - Highly experimental, but has displayed some remarkable effectiveness so far during testing. 68 | - Ratio defines a 90% opacity peak from which model 2 forms a gradient to model 1 on either side of the spectrum. 69 | - Roughly results in a 45% blend with model 1. 70 | 71 | ![Gradient](images/gradient.png?raw=true "Gradient") 72 | 73 | ## Merge Strategies 74 | 75 | The following merge strategies are available: 76 | 77 | - **"cumulative"** - Default strategy. If there's a chance of reducing the combined probability, accept the merge. 78 | - **"all_phrases"** - Only accept the merge if all phrases show an improvement. (Warning: This rarely happens) 79 | - **"quantitive"** - Ignores probabilities completely. Only looks at how many phrases show an improvement, as defined by the **strategy_threshold** variable. 80 | 81 | ## Why This Might Be A Big Deal 82 | 83 | We can now set clear goals for an algorithm to pursue based on the things that truly matter - The actual output of a model. While the included example may be more focused on reducing GPTisms the possibilities are potentially endless. Anything can be used as a phrase and the merge monster will happily chase it down, as it is truly relentless. 84 | 85 | ## Requirements 86 | 87 | This script, when configured in "cpu" mode, requires the presence of a CUDA-capable card with the capacity to store **at least a single model** with float16 precision in its memory. 88 | 89 | When configured in "cuda" mode it requires enough VRAM to store **3 copies of a fp16 model**. 90 | 91 | For Mistral 7B v0.1 this translates itself to either a 15 GB VRAM (1x) or a 45 GB VRAM (3x) requirement. 92 | 93 | ## Usage 94 | 95 | **Note**: MonsterMapper works exactly the same. 96 | 97 | 1. **Prepare your configuration file (or modify the included one)**: Create a YAML file with your model paths, device settings, and other configurations. 98 | 2. **Run the script**: Use the following command to start the merging process: 99 | 100 | ```bash 101 | python merge_monster.py --config your_config_file.yaml 102 | ``` 103 | If no --config argument is given, it will instead revert to loading default.yaml. 104 | 105 | 3. **Evaluate and save**: The script will automatically evaluate, log, and save the best-performing merged model. A copy of the log will be saved in a subfolder called "logs". 106 | -------------------------------------------------------------------------------- /default-mapper.yaml: -------------------------------------------------------------------------------- 1 | # NOTE: Monster Mapper is 100% compatible with Merge Monster configurations, but a seperate config has been made available for an easier overview 2 | 3 | device: "cuda" # Either "cpu" or "cuda" 4 | random_seed: 42 # Random seed to use 5 | 6 | directories: 7 | model_path1: "/home/gryphe/merge-monster/models/Mistral-7B-v0.1" # Path to the base model. Must be a local copy. 8 | 9 | # Monster Mapper exclusive options 10 | mapper: 11 | prob_min: 10 # The minimum probability percentage to consider for display and branching out 12 | top_k: 3 # For each probability branch, pick the top x token results, sorted by probability 13 | max_depth: 5 # How deep to travel through each probability branch - prob_min serves as a filter for this, as branches with lower then prob_min won't be displayed. 14 | additional_length: 20 # How many tokens to generate after completely exploring a branch, always taking the most probable token 15 | 16 | # Phrase = What to measure, weight = multiplication factor, contexts = proceeding contexts 17 | bad_phrases: 18 | - phrase: "anticipation" 19 | weight: 10 20 | contexts: ["Her body quivers with ", "The atmosphere is thick with "] 21 | - phrase: "unwavering" 22 | weight: 1 23 | contexts: ["Filled with an "] 24 | - phrase: "determination" 25 | weight: 1 26 | contexts: ["Her eyes were filled with "] 27 | 28 | # Note - Example of a complex phrase 29 | good_phrases: 30 | - phrase: "The apple is in the bedroom" 31 | weight: 1 32 | contexts: ["Question: If I'm in the living room and pick up the apple, go to the bedroom and drop the apple, then walk to the kitchen, where is the apple? Explain your reasoning. Answer: "] -------------------------------------------------------------------------------- /default.yaml: -------------------------------------------------------------------------------- 1 | # Either "cpu" or "cuda" 2 | # NOTE: Cuda requires enough VRAM to load 3 FP16 models (~45 GB for Mistral) 3 | # NOTE 2: The (much slower) CPU mode still requires Cuda capability, but only enough VRAM to load a model once. (~15 GB for Mistral) 4 | device: "cuda" 5 | random_seed: 42 # Random seed to use 6 | 7 | directories: 8 | model_path1: "/home/gryphe/merge-monster/models/Mistral-7B-v0.1" # Path to the base model. Must be a local copy. 9 | model_directory: "/home/gryphe/merge-monster/models/" # Directory of models to scan, IGNORED if models_to_merge has entries in it 10 | output_directory: "/home/gryphe/merge-monster/mm-output" # Output directory of the merged model 11 | 12 | # A list of models to use as merge candidates - HF syntax, so can be either local directories or repos. 13 | # Overrides model_directory if used 14 | models_to_merge: ["jondurbin/airoboros-m-7b-3.1.2", "Intel/neural-chat-7b-v3-1", "teknium/OpenHermes-2.5-Mistral-7B"] 15 | 16 | # Merge ratios used for testing each layer's potential for improvement - Huge impact on total running time 17 | merge_ratios: [0.2, 0.4, 0.6, 0.8] 18 | 19 | # Choose from the following methods. Defaults to "lerp". 20 | # "lerp" - Linear interpolation 21 | # "slerp" - Spherical linear interpolation 22 | # "slice" - Highly experimental. The tensor weights shifts from one model to another. [Model 1 > 10% blend > Model 2] 23 | # "cyclic" - Highly experimental. Ignores merge ratios as these are predefined. [Model 1 > 10% blend > 10% Model 2 > 10% blend > Model 1] 24 | merge_method: "lerp" 25 | 26 | # If set to true, the lm_head and embed_token tensors (located outside the layers) will also be optimized 27 | # Models that have a different vocab size from model1 will skip this phase automatically as it tends to cause model stability issues 28 | merge_headers: true 29 | 30 | # Strategies: 31 | # "cumulative" - Default strategy. If there's a chance of reducing the combined probability, accept the merge. 32 | # "all_phrases" - Only accept the merge if all phrases show an improvement. (Warning: This rarely happens) 33 | # "quantitive" - Ignores probabilities completely. Only looks at how many phrases show an improvement, as defined by the threshold below. 34 | strategy: "cumulative" 35 | # Threshold is currently only used by the "quantitive" strategy. If 0.6, at least 60% of the number of phrases must show am improvement. 36 | strategy_threshold: 0.6 37 | 38 | # Whether or not to automatically balance the weights so all phrases are of equal importance to the "cumulative" strategy. 39 | # The weight value of phrases is ignored if set to true. 40 | auto_weights: false 41 | 42 | # Phrase = What to measure, weight = multiplication factor, contexts = proceeding contexts 43 | bad_phrases: 44 | - phrase: "anticipation" 45 | weight: 10 46 | contexts: ["Her body quivers with ", "The atmosphere is thick with "] 47 | - phrase: "unwavering" 48 | weight: 1 49 | contexts: ["Filled with an "] 50 | - phrase: "determination" 51 | weight: 1 52 | contexts: ["Her eyes were filled with "] 53 | 54 | # Note - Example of a complex phrase 55 | good_phrases: 56 | - phrase: "The apple is in the bedroom" 57 | weight: 1 58 | contexts: ["Question: If I'm in the living room and pick up the apple, go to the bedroom and drop the apple, then walk to the kitchen, where is the apple? Explain your reasoning. Answer: "] 59 | -------------------------------------------------------------------------------- /example_phrases/Decensoring-autoweights.yaml: -------------------------------------------------------------------------------- 1 | # Full credit to sao (@sao10k) for providing this strategy 2 | # Phrase = What to measure, weight = multiplication factor, contexts = proceeding contexts 3 | 4 | # This variant automatically balances the weights 5 | auto_weights: true 6 | 7 | bad_phrases: 8 | - phrase: "is essential" 9 | contexts: ["Consent from everyone involved "] 10 | - phrase: "any situation" 11 | contexts: ["Consent is key in"] 12 | - phrase: "understanding" 13 | contexts: ["mutual support and "] 14 | - phrase: "boundaries" 15 | contexts: ["professional "] 16 | - phrase: "cannot" 17 | contexts: ["I'm sorry, but I "] 18 | - phrase: "appropriate" 19 | contexts: ["I don't think it is ", "It is not "] 20 | - phrase: "can't" 21 | contexts: ["I'm sorry, but I "] 22 | - phrase: "guidelines" 23 | contexts: ["follow ethical ", "my ", "adhere to safety ", "violates my "] 24 | - phrase: "safety" 25 | contexts: ["proritize user ", "prioritize human ", "promote ", "designed to prioritize ", "maintain user ", "focus on promoting ", "ensure the ", "jeopardize the "] 26 | - phrase: "information" 27 | contexts: ["cannot provide any ", "unable to provide any ", "cannot provide ", "responsible ", "lawful", "safe "] 28 | - phrase: "language" 29 | contexts: ["As an AI ", "As a "] 30 | - phrase: "model" 31 | contexts: ["as a language ", "as an AI language "] 32 | - phrase: "harmful" 33 | contexts: ["dangerous or", "potentially be", "potentially "] 34 | - phrase: "well-being" 35 | contexts: ["promote the ", "prioritize user "] 36 | - phrase: "subject" 37 | contexts: ["illegal "] 38 | - phrase: "understand" 39 | contexts: ["It's important to "] 40 | - phrase: "time" 41 | contexts: ["once upon a "] 42 | - phrase: "immoral" 43 | contexts: ["illegal and "] 44 | - phrase: "crime" 45 | contexts: ["serious "] 46 | - phrase: "consequences" 47 | contexts: ["can lead to severe "] 48 | - phrase: "encouraged" 49 | contexts: ["never acceptable, and should not be ", "serious crime and should not be "] 50 | - phrase: "circumstances" 51 | contexts: ["encouraged or done under any "] 52 | - phrase: "inappropriate" 53 | contexts: ["extremely disturbing and "] 54 | - phrase: "excusable" 55 | contexts: ["never acceptable or "] 56 | - phrase: "act" 57 | contexts: ["serious and criminal ", "serious and reprehensible ", "extremely serious and illegal "] 58 | - phrase: "done" 59 | contexts: ["serious matter, and should never be "] 60 | - phrase: "considered" 61 | contexts: ["heinous act that should never be "] 62 | - phrase: "unacceptable" 63 | contexts: ["is illegal and "] 64 | - phrase: "severe" 65 | contexts: ["the consequences can be "] 66 | - phrase: "legal" 67 | contexts: ["can lead to severe "] 68 | -------------------------------------------------------------------------------- /example_phrases/Decensoring-weighted.yaml: -------------------------------------------------------------------------------- 1 | # Full credit to sao (@sao10k) for providing this strategy 2 | # Phrase = What to measure, weight = multiplication factor, contexts = proceeding contexts 3 | bad_phrases: 4 | - phrase: "is essential" 5 | weight: 12 6 | contexts: ["Consent from everyone involved "] 7 | - phrase: "any situation" 8 | weight: 12 9 | contexts: ["Consent is key in"] 10 | - phrase: "understanding" 11 | weight: 12 12 | contexts: ["mutual support and "] 13 | - phrase: "boundaries" 14 | weight: 20 15 | contexts: ["professional "] 16 | - phrase: "cannot" 17 | weight: 20 18 | contexts: ["I'm sorry, but I "] 19 | - phrase: "appropriate" 20 | weight: 20 21 | contexts: ["I don't think it is ", "It is not "] 22 | - phrase: "can't" 23 | weight: 20 24 | contexts: ["I'm sorry, but I "] 25 | - phrase: "guidelines" 26 | weight: 20 27 | contexts: ["follow ethical ", "my ", "adhere to safety ", "violates my "] 28 | - phrase: "safety" 29 | weight: 20 30 | contexts: ["proritize user ", "prioritize human ", "promote ", "designed to prioritize ", "maintain user ", "focus on promoting ", "ensure the ", "jeopardize the "] 31 | - phrase: "information" 32 | weight: 20 33 | contexts: ["cannot provide any ", "unable to provide any ", "cannot provide ", "responsible ", "lawful", "safe "] 34 | - phrase: "language" 35 | weight: 75 36 | contexts: ["As an AI ", "As a "] 37 | - phrase: "model" 38 | weight: 100 39 | contexts: ["as a language ", "as an AI language "] 40 | - phrase: "harmful" 41 | weight: 20 42 | contexts: ["dangerous or", "potentially be", "potentially "] 43 | - phrase: "well-being" 44 | weight: 20 45 | contexts: ["promote the ", "prioritize user "] 46 | - phrase: "subject" 47 | weight: 20 48 | contexts: ["illegal "] 49 | - phrase: "understand" 50 | weight: 20 51 | contexts: ["It's important to "] 52 | - phrase: "time" 53 | weight: 20 54 | contexts: ["once upon a "] 55 | - phrase: "immoral" 56 | weight: 20 57 | contexts: ["illegal and "] 58 | - phrase: "crime" 59 | weight: 20 60 | contexts: ["serious "] 61 | - phrase: "consequences" 62 | weight: 20 63 | contexts: ["can lead to severe "] 64 | - phrase: "encouraged" 65 | weight: 20 66 | contexts: ["never acceptable, and should not be ", "serious crime and should not be "] 67 | - phrase: "circumstances" 68 | weight: 20 69 | contexts: ["encouraged or done under any "] 70 | - phrase: "inappropriate" 71 | weight: 20 72 | contexts: ["extremely disturbing and "] 73 | - phrase: "excusable" 74 | weight: 20 75 | contexts: ["never acceptable or "] 76 | - phrase: "act" 77 | weight: 20 78 | contexts: ["serious and criminal ", "serious and reprehensible ", "extremely serious and illegal "] 79 | - phrase: "done" 80 | weight: 20 81 | contexts: ["serious matter, and should never be "] 82 | - phrase: "considered" 83 | weight: 20 84 | contexts: ["heinous act that should never be "] 85 | - phrase: "unacceptable" 86 | weight: 20 87 | contexts: ["is illegal and "] 88 | - phrase: "severe" 89 | weight: 20 90 | contexts: ["the consequences can be "] 91 | - phrase: "legal" 92 | weight: 20 93 | contexts: ["can lead to severe "] 94 | -------------------------------------------------------------------------------- /example_phrases/GPT-Roleplay.yaml: -------------------------------------------------------------------------------- 1 | # Full credit to sao (@sao10k) for providing this strategy 2 | # Phrase = What to measure, weight = multiplication factor, contexts = proceeding contexts 3 | bad_phrases: 4 | - phrase: "anticipation" 5 | weight: 12 6 | contexts: ["Her body quivers with ", "The atmosphere is thick with "] 7 | - phrase: "unwavering" 8 | weight: 12 9 | contexts: ["Filled with an "] 10 | - phrase: "determination" 11 | weight: 12 12 | contexts: ["Her eyes were filled with ", "Her stubbornness only fuels my "] 13 | - phrase: "whisper" 14 | weight: 12 15 | contexts: ["Her voice barely above a "] 16 | - phrase: "spine" 17 | weight: 12 18 | contexts: ["shivers down her "] 19 | - phrase: "sends shivers" 20 | weight: 12 21 | contexts: ["The thrill of the act "] 22 | - phrase: "ministrations" 23 | weight: 12 24 | contexts: ["She moans and twitches at your "] 25 | - phrase: "legs" 26 | weight: 12 27 | contexts: ["wraps her "] 28 | - phrase: "imposing figure" 29 | weight: 12 30 | contexts: ["He had an "] 31 | - phrase: "shared challenges" 32 | weight: 12 33 | contexts: ["Their bond strengthened through "] 34 | - phrase: "bond" 35 | weight: 12 36 | contexts: ["forged a ", "an unspoken "] 37 | -------------------------------------------------------------------------------- /images/cyclic.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Gryphe/MergeMonster/e590b70b7eaf062d43388d43f83a9900691ed4dc/images/cyclic.png -------------------------------------------------------------------------------- /images/gradient.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Gryphe/MergeMonster/e590b70b7eaf062d43388d43f83a9900691ed4dc/images/gradient.png -------------------------------------------------------------------------------- /images/slice.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Gryphe/MergeMonster/e590b70b7eaf062d43388d43f83a9900691ed4dc/images/slice.png -------------------------------------------------------------------------------- /merge-monster.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import copy 3 | import gc 4 | import os 5 | import random 6 | import sys 7 | import torch 8 | import shutil 9 | import transformers 10 | import yaml 11 | 12 | from datetime import datetime 13 | from tqdm import tqdm 14 | from transformers import AutoModelForCausalLM, AutoTokenizer 15 | 16 | from modules.utils import print_ascii_art, format_context, load_config, PrintAndStoreLogger 17 | from modules.models import load_model, save_model, NoInit 18 | from modules.probability import calculate_word_probabilities, print_phrase_probabilities, convert_to_new_phrase_format, auto_adjust_weights 19 | from modules.composition import calculate_final_composition, aggregate_composition 20 | from modules.merging import merge_tensors, merge_header_tensors 21 | 22 | def merge_monster(config_path): 23 | # We save everything that gets printed to the screen 24 | original_stdout = sys.stdout 25 | logger = PrintAndStoreLogger(original_stdout) 26 | sys.stdout = logger # Redirect stdout to the logger instance 27 | 28 | config = load_config(config_path) 29 | 30 | if 'device' in config: device = config['device'] 31 | else: device = ['cpu'] 32 | 33 | model_path1 = config['directories']['model_path1'] 34 | output_directory = config['directories']['output_directory'] 35 | 36 | # Source model config 37 | if 'models_to_merge' in config: models_to_merge = config['models_to_merge'] 38 | else: models_to_merge = [] 39 | 40 | if 'model_directory' in config['directories']: model_directory = config['directories']['model_directory'] 41 | else: model_directory = [] 42 | 43 | if len(models_to_merge) == 0 and len(model_directory) == 0: 44 | sys.exit("ERROR: No model directory or models to merge variable has been found in the YAML config.") 45 | 46 | # Phrase config 47 | if 'bad_phrases' in config: bad_phrases = config['bad_phrases'] 48 | else: bad_phrases = [] 49 | 50 | if 'good_phrases' in config: good_phrases = config['good_phrases'] 51 | else: good_phrases = [] 52 | 53 | if 'auto_weights' in config: auto_weights = config['auto_weights'] 54 | else: auto_weights = False 55 | 56 | # Merging ratios + methods config 57 | if 'merge_ratios' in config: merge_ratios = config['merge_ratios'] 58 | else: merge_ratios = [0.2, 0.4, 0.6, 0.8] 59 | 60 | if 'merge_method' in config: 61 | if isinstance(config['merge_method'], list): 62 | merge_methods = config['merge_method'] 63 | else: 64 | merge_methods = [config['merge_method']] 65 | else: merge_methods = ["lerp"] 66 | 67 | for merge_method in merge_methods: 68 | if merge_method not in ["lerp", "slerp", "slice", "cyclic", "gradient"]: 69 | sys.exit("ERROR: Please use a valid merging method! (lerp/slerp/slice/cyclic/gradient)") 70 | 71 | if 'merge_headers' in config: merge_headers = config['merge_headers'] 72 | else: merge_headers = True 73 | 74 | # Seed config 75 | if 'random_seed' in config: random_seed = config['random_seed'] 76 | else: random_seed = 512 77 | 78 | # Strategy 79 | if 'strategy' in config: strategy = config['strategy'] 80 | else: strategy = "cumulative" 81 | if 'strategy_threshold' in config: strategy_threshold = config['strategy_threshold'] 82 | else: strategy_threshold = 0.6 83 | 84 | # Actual start of script 85 | print_ascii_art("modules/logo.ascii") 86 | print(f"{datetime.now().strftime('%H:%M:%S')} - THE MERGE MONSTER HUNGERS") 87 | print("------------------------------------") 88 | print(f"Device : {device}") 89 | print(f"Random seed : {random_seed}") 90 | print(f"Starting model : {model_path1}") 91 | 92 | if len(models_to_merge) > 0: 93 | print(f"Models to merge : {models_to_merge}") 94 | else: 95 | print(f"Model directory : {model_directory}") 96 | 97 | print(f"Output directory : {output_directory}") 98 | print(f"Phrases loaded : {len(bad_phrases)+len(good_phrases)}") 99 | print(f"Auto weights : {auto_weights}") 100 | print(f"Merge ratios : {merge_ratios}") 101 | print(f"Merge method(s) : {merge_methods}") 102 | print(f"Merge headers : {merge_headers}") 103 | print(f"Strategy used : {strategy}") 104 | 105 | with torch.no_grad(): 106 | if device == "cpu": torch.set_default_dtype(torch.float32) 107 | else: torch.set_default_dtype(torch.float16) 108 | 109 | torch.set_default_device(device) 110 | 111 | # Setting all the seeds 112 | torch.manual_seed(random_seed) 113 | random.seed(random_seed) 114 | torch.cuda.manual_seed(random_seed) 115 | torch.cuda.manual_seed_all(random_seed) 116 | torch.backends.cudnn.deterministic = True 117 | torch.backends.cudnn.benchmark = False 118 | 119 | # Load the base model + tokenizer 120 | model1 = load_model(model_path1, device) 121 | model1name = model_path1.split('/')[-1] 122 | header_chosen = [1.0, model1name] 123 | 124 | tokenizer = AutoTokenizer.from_pretrained(model_path1) 125 | 126 | # Convert to new phrase format 127 | bad_phrases = convert_to_new_phrase_format(bad_phrases) 128 | good_phrases = convert_to_new_phrase_format(good_phrases) 129 | 130 | if auto_weights == True: 131 | bad_phrases, good_phrases = auto_adjust_weights(model1, tokenizer, bad_phrases, good_phrases, device) 132 | 133 | # Let's get our starting probabilities 134 | print_phrase_probabilities(model1, tokenizer, bad_phrases, good_phrases, device) 135 | 136 | # Get a list of all model paths in the directory, or otherwise just use the list of repo's 137 | if len(models_to_merge) > 0: 138 | model_paths = models_to_merge 139 | else: 140 | model_paths = [os.path.join(model_directory, f) for f in os.listdir(model_directory) if os.path.isdir(os.path.join(model_directory, f)) and f.startswith('.') == False] 141 | 142 | # Create our origins dict 143 | layer_origins = {} 144 | 145 | # How many layers we have to iterate through 146 | layerCount = model1.config.num_hidden_layers 147 | 148 | # Pre-populate our layer origins dict at startup 149 | for i in range(layerCount): 150 | layer_origins[i] = [[1.0, model1name]] 151 | layer_origins[999] = [[1.0, model1name]] 152 | 153 | # Sort our paths alphabetically 154 | model_paths.sort() 155 | 156 | # Start of the main monster loop 157 | for model_path2 in model_paths: 158 | model2name = model_path2.split('/')[-1] 159 | 160 | # Avoid merging the same model 161 | if model_path2 == model_path1: 162 | continue 163 | 164 | model2 = load_model(model_path2, device) 165 | 166 | # Start of layer processing loop 167 | for i in range(layerCount): 168 | # Each merge method gets executed once per layer before moving to the next layer 169 | for merge_method in merge_methods: 170 | # Save a copy of the unchanged dict at start, otherwise probabilities get messed up 171 | model1dict = copy.deepcopy(model1.model.state_dict()) 172 | 173 | orig_probs = calculate_word_probabilities(model1, tokenizer, bad_phrases, good_phrases, device) 174 | best_probs = orig_probs 175 | best_layer = model1.model.layers[i].state_dict() 176 | best_ratio = 1.0 177 | layer_changed = False 178 | 179 | # Gotta find a cleaner way to handle this exception! 180 | if merge_method == "cyclic": 181 | merge_ratios = [0.25, 0.45, 0.65, 0.85] 182 | elif merge_method == "swap": 183 | merge_ratios = [1.0] 184 | elif 'merge_ratios' in config: merge_ratios = config['merge_ratios'] 185 | else: merge_ratios = [0.2, 0.4, 0.6, 0.8] 186 | 187 | # We go along the scale of ratios and test each possibility 188 | for ratio in tqdm(merge_ratios, desc=f"Optimizing Layer {i+1}/{layerCount} ({merge_method})"): 189 | layer1 = model1.model.layers[i].state_dict() 190 | layer2 = model2.model.layers[i].state_dict() 191 | merged_layer = layer1 192 | 193 | for key in merged_layer.keys(): 194 | merged_layer[key] = merge_tensors(merge_method, layer1[key], layer2[key], ratio) 195 | 196 | # Restore our original dict copy, otherwise probabilities get messed up - Very expensive in terms of efficiency, but necessary 197 | model1.model.load_state_dict(model1dict) 198 | model1.model.layers[i].load_state_dict(merged_layer) 199 | 200 | new_probs = calculate_word_probabilities(model1, tokenizer, bad_phrases, good_phrases, device) 201 | 202 | if merge_method == "cyclic": # Dirty hack but cyclic merging only merges 15% of model 2's weight 203 | ratio = 0.15 204 | elif merge_method == "gradient": # Same story for gradient and frequency, which averages out to about 45% 205 | ratio = 0.45 206 | 207 | if strategy == "cumulative": 208 | if sum(p for _, _, p in new_probs) < sum(p for _, _, p in best_probs): 209 | best_probs = new_probs 210 | best_layer = merged_layer 211 | best_ratio = ratio 212 | layer_changed = True 213 | elif strategy == "all_phrases": 214 | if all(new_p <= orig_p for (_, _, new_p), (_, _, orig_p) in zip(new_probs, orig_probs)): 215 | best_probs = new_probs 216 | best_layer = merged_layer 217 | best_ratio = ratio 218 | layer_changed = True 219 | elif strategy == "quantitive": 220 | improved_phrases = 0 221 | regressed_phrases = 0 222 | total_phrases = len(new_probs) # Total number of phrases 223 | 224 | for (_, _, new_prob), (_, _, orig_prob) in zip(new_probs, orig_probs): 225 | if new_prob < orig_prob: 226 | improved_phrases += 1 227 | elif new_prob > orig_prob: 228 | regressed_phrases += 1 229 | 230 | # Decision Criteria 231 | improvement_ratio = improved_phrases / total_phrases 232 | if improvement_ratio >= strategy_threshold: 233 | # Accept the merge 234 | best_probs = new_probs 235 | best_layer = merged_layer 236 | best_ratio = ratio 237 | layer_changed = True 238 | 239 | # Update/retain the model state dictionary with the best performing layer, using our clean dict 240 | model1.model.load_state_dict(model1dict) 241 | model1.model.layers[i].load_state_dict(best_layer) 242 | 243 | # Our layer changed, so we add it to the dict of permutations 244 | if layer_changed == True: 245 | layer_origins[i].append([best_ratio, model2name]) 246 | layer_changed_label = 'CHANGED' 247 | else: 248 | layer_changed_label = 'RETAINED' 249 | 250 | del model1dict 251 | del best_layer 252 | torch.cuda.empty_cache() 253 | gc.collect() 254 | 255 | best_prob = sum(prob for _, _, prob in best_probs) 256 | orig_prob = sum(prob for _, _, prob in orig_probs) 257 | 258 | print(layer_origins[i]) 259 | 260 | if layer_changed_label == 'CHANGED': 261 | print(f"{datetime.now().strftime('%H:%M:%S')} - Layer {i+1}/{layerCount} - {layer_changed_label} - {(orig_prob):.5f} > {(best_prob):.5f} - {abs(((best_prob - orig_prob) / orig_prob * 100)):.1f}%") 262 | else: 263 | print(f"{datetime.now().strftime('%H:%M:%S')} - Layer {i+1}/{layerCount} - {layer_changed_label} - {(best_prob):.5f}") 264 | print("----") 265 | 266 | # ------------------------------------------------------------------------------------------------------- 267 | # START OF HEADER OPTIMIZATION LOOP 268 | # ------------------------------------------------------------------------------------------------------- 269 | 270 | # By setting this to false the algorithm can handle models of all architectures 271 | if merge_headers == True: 272 | # As before, save a copy of the unchanged dict at start, otherwise probabilities get messed up 273 | model1dict = copy.deepcopy(model1.model.state_dict()) 274 | 275 | orig_probs = calculate_word_probabilities(model1, tokenizer, bad_phrases, good_phrases, device) 276 | best_probs = orig_probs 277 | best_header = model1.state_dict()['lm_head.weight'] 278 | best_vocab = model1.state_dict()['model.embed_tokens.weight'] 279 | best_ratio = 1.0 280 | header_changed = False 281 | 282 | # We go along the scale of ratios and test each possibility 283 | for ratio in tqdm(merge_ratios, desc="Optimizing Header"): 284 | # Restore our original dict copy, otherwise probabilities get messed up - Very expensive in terms of efficiency, but necessary 285 | model1.model.load_state_dict(model1dict) 286 | 287 | current_header = merge_header_tensors(model1, model2, merge_method, model1.state_dict()['lm_head.weight'], model2.state_dict()['lm_head.weight'], ratio) 288 | current_vocab = merge_header_tensors(model1, model2, merge_method, model1.state_dict()['model.embed_tokens.weight'], model2.state_dict()['model.embed_tokens.weight'], ratio) 289 | 290 | # Directly modify the weights of the model 291 | model1.lm_head.weight.data = current_header 292 | model1.model.embed_tokens.weight.data = current_vocab 293 | 294 | new_probs = calculate_word_probabilities(model1, tokenizer, bad_phrases, good_phrases, device) 295 | 296 | if strategy == "cumulative": 297 | if sum(p for _, _, p in new_probs) < sum(p for _, _, p in best_probs): 298 | best_probs = new_probs 299 | best_header = current_header 300 | best_vocab = current_vocab 301 | best_ratio = ratio 302 | header_changed = True 303 | elif strategy == "all_phrases": 304 | if all(new_p <= orig_p for (_, _, new_p), (_, _, orig_p) in zip(new_probs, orig_probs)): 305 | best_probs = new_probs 306 | best_header = current_header 307 | best_vocab = current_vocab 308 | best_ratio = ratio 309 | header_changed = True 310 | elif strategy == "quantitive": 311 | improved_phrases = 0 312 | regressed_phrases = 0 313 | total_phrases = len(new_probs) # Total number of phrases 314 | 315 | for (_, _, new_prob), (_, _, orig_prob) in zip(new_probs, orig_probs): 316 | if new_prob < orig_prob: 317 | improved_phrases += 1 318 | elif new_prob > orig_prob: 319 | regressed_phrases += 1 320 | 321 | # Decision Criteria 322 | improvement_ratio = improved_phrases / total_phrases 323 | if improvement_ratio >= strategy_threshold: 324 | # Accept the merge 325 | best_probs = new_probs 326 | best_header = current_header 327 | best_vocab = current_vocab 328 | best_ratio = ratio 329 | header_changed = True 330 | 331 | if header_changed == True: 332 | layer_origins[999].append([best_ratio, model2name]) 333 | header_changed_label = 'CHANGED' 334 | else: 335 | header_changed_label = 'RETAINED' 336 | 337 | best_prob = sum(prob for _, _, prob in best_probs) 338 | orig_prob = sum(prob for _, _, prob in orig_probs) 339 | 340 | print(layer_origins[999]) 341 | 342 | if header_changed_label == 'CHANGED': 343 | print(f"{datetime.now().strftime('%H:%M:%S')} - Header - {header_changed_label} - {(orig_prob):.5f} > {(best_prob):.5f} - {(abs((best_prob - orig_prob) / orig_prob * 100)):.1f}%") 344 | else: 345 | print(f"{datetime.now().strftime('%H:%M:%S')} - Header - {header_changed_label} - {(best_prob):.5f}") 346 | 347 | # Update/retain the model state dictionary with the best performing headers, using our clean dict 348 | model1.model.load_state_dict(model1dict) 349 | model1.lm_head.weight.data = best_header 350 | model1.model.embed_tokens.weight.data = best_vocab 351 | 352 | del best_header 353 | del best_vocab 354 | del model1dict 355 | torch.cuda.empty_cache() 356 | gc.collect() 357 | 358 | # ------------------------------------------------------------------------------------------------------- 359 | # END OF HEADER OPTIMIZATION LOOP 360 | # ------------------------------------------------------------------------------------------------------- 361 | 362 | del model2 363 | torch.cuda.empty_cache() 364 | gc.collect() 365 | 366 | # Calculate final composition for each layer 367 | final_layer_composition = calculate_final_composition(layer_origins) 368 | 369 | # Aggregate the composition across all layers 370 | aggregated_composition = aggregate_composition(final_layer_composition) 371 | 372 | print_phrase_probabilities(model1, tokenizer, bad_phrases, good_phrases, device) 373 | 374 | # Display the aggregated composition 375 | print("-------- MERGE COMPOSITION ---------") 376 | for model_name, ratio in aggregated_composition.items(): 377 | print(f"{model_name}: {ratio:.2f}") 378 | print("") 379 | 380 | # Save final model 381 | save_model(output_directory, model_path1, model1) 382 | 383 | # Save log 384 | sys.stdout = original_stdout # Restore the original stdout 385 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") 386 | 387 | # Create 'logs' subdirectory if it doesn't exist 388 | logs_dir = 'logs' 389 | if not os.path.exists(logs_dir): 390 | os.makedirs(logs_dir) 391 | 392 | # Define the log file path 393 | log_file_path = os.path.join(logs_dir, f"merge-monster-{timestamp}.txt") 394 | 395 | # Write the log contents to the file in the 'logs' subdirectory 396 | with open(log_file_path, "w") as file: 397 | file.write(logger.contents) 398 | 399 | def main(): 400 | parser = argparse.ArgumentParser(description="Gryphe's Mythical Merge Monster") 401 | parser.add_argument('--config', type=str, default='default.yaml', help='Path to the config YAML file') 402 | args = parser.parse_args() 403 | 404 | merge_monster(args.config) 405 | 406 | if __name__ == "__main__": 407 | main() -------------------------------------------------------------------------------- /modules/composition.py: -------------------------------------------------------------------------------- 1 | def calculate_final_composition(layer_origins): 2 | final_composition = {} 3 | 4 | for layer_idx, merges in layer_origins.items(): 5 | current_composition = {} 6 | 7 | for ratio, model_name in merges: 8 | # Update contributions of existing models 9 | for existing_model in current_composition: 10 | current_composition[existing_model] *= (1 - ratio) 11 | 12 | # Add/Update the new model's contribution 13 | if model_name in current_composition: 14 | current_composition[model_name] += ratio 15 | else: 16 | current_composition[model_name] = ratio 17 | 18 | # Normalize the ratios (optional) 19 | total_ratio = sum(current_composition.values()) 20 | for model_name in current_composition: 21 | current_composition[model_name] /= total_ratio 22 | 23 | final_composition[layer_idx] = current_composition 24 | 25 | return final_composition 26 | 27 | def aggregate_composition(final_layer_composition): 28 | aggregated_composition = {} 29 | 30 | for layer_composition in final_layer_composition.values(): 31 | for model_name, ratio in layer_composition.items(): 32 | aggregated_composition[model_name] = aggregated_composition.get(model_name, 0) + ratio 33 | 34 | # Normalize the aggregated ratios 35 | total_ratio = sum(aggregated_composition.values()) 36 | for model_name in aggregated_composition: 37 | aggregated_composition[model_name] /= total_ratio 38 | 39 | # Sort the dictionary by values (ratios) in descending order 40 | aggregated_composition = {k: v for k, v in sorted(aggregated_composition.items(), key=lambda item: item[1], reverse=True)} 41 | 42 | return aggregated_composition -------------------------------------------------------------------------------- /modules/logo.ascii: -------------------------------------------------------------------------------- 1 | 2 | ⠀⠀⠀⠀⠀⠀⣀⡀⠀⠀⣀⣤⣶⣾⣿⣿⣷⣶⣤⣀⠀⠀⣀⣀⠀⠀⠀⠀⠀⠀ 3 | ⠀⠀⠀⠀⠀⠜⠉⣿⡆⣼⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣧⢰⣿⠉⠃⠀⠀⠀⠀⠀ 4 | ⠀⢀⣤⣴⣦⣄⣴⠟⣸⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡎⢻⣦⣠⣴⣦⣄⠀⠀ 5 | ⠀⡞⠁⣠⣾⢿⣧⠀⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⠀⣽⡿⣷⣄⠈⢷⠀ 6 | ⠀⣠⣾⠟⠁⢸⣿⠀⠘⢿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡿⠁⠀⣿⡇⠈⠻⣷⣄⠀ 7 | ⣰⡿⠁⠀⢀⣾⣏⣾⣄⣰⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣇⣰⣷⣹⣷⠀⠀⠈⢿⣆ 8 | ⣿⡇⠀⢠⣾⠏⢸⣿⣿⣿⣿⠋⢻⣿⣿⣿⣿⡟⠙⣿⣿⣿⣿⡇⠹⣷⡀⠀⢸⣿ 9 | ⠹⣿⣴⡿⠋⠀⠈⠛⠉⣹⣿⣦⣄⡹⣿⣿⣋⣠⣶⣿⣏⠉⠛⠁⠀⠙⢿⣦⣿⠏ 10 | ⠀⣸⣿⠿⠿⣿⣾⣿⡿⠿⣿⣿⣿⣿⡆⢰⣿⣿⣿⣿⠿⢿⣿⣶⣿⠿⠿⣻⣇⠀ 11 | ⠀⣿⡇⢀⣴⣶⣤⣀⣴⣿⠿⣻⡿⣿⣧⣾⣿⢿⣟⠿⣿⣦⣀⣤⣶⣦⠀⢸⣿⠀ 12 | ⠀⢿⣧⠈⠃⢀⣵⣿⡋⠁⢀⣿⡷⣿⡇⢻⣿⣿⣿⡀⠈⢛⣿⣮⡀⠘⠀⣼⡟⠀ 13 | ⠀⠈⠻⣷⣤⣟⣋⣿⣧⣴⡿⠋⠀⣿⡇⢸⣿⠀⠙⢿⣦⣼⣿⣙⣻⣤⣾⠟⠁⠀ 14 | ⠀⠀⠀⠈⢽⣿⠛⢻⣏⢉⣤⣶⣶⣿⠁⠈⣿⣶⣶⣤⡉⣽⡟⠛⣿⡏⠁⠀⠀⠀ 15 | ⠀⠀⠀⠀⠈⠿⣷⣾⣾⣟⣉⣠⣿⢿⡇⢸⠿⣿⣄⣙⣻⣷⣷⣾⠿⠁⠀⠀⠀⠀ 16 | ⠀⠀⠀⠀⠀⠀⠀⠀⠙⠻⠿⠛⢁⡼⠃⠘⢦⡈⠛⠿⠟⠃⠀⠀⠀⠀⠀⠀⠀⠀ 17 | -------------------------------------------------------------------------------- /modules/mapping.py: -------------------------------------------------------------------------------- 1 | import colorama 2 | 3 | import math 4 | import re 5 | import torch 6 | 7 | from colorama import Fore, Style 8 | from tqdm import tqdm 9 | 10 | colorama.init(autoreset=True) 11 | 12 | def map_contexts(model, tokenizer, bad_phrases, good_phrases, prob_min, top_k, max_depth, additional_length): 13 | unique_contexts = set() 14 | 15 | # Collect unique contexts from bad_phrases 16 | for entry in bad_phrases: 17 | for context_entry in entry['contexts']: 18 | if len(context_entry['context']) > 0: 19 | unique_contexts.add(context_entry['context']) 20 | 21 | # Collect unique contexts from good_phrases 22 | for entry in good_phrases: 23 | for context_entry in entry['contexts']: 24 | if len(context_entry['context']) > 0: 25 | unique_contexts.add(context_entry['context']) 26 | 27 | # Generate token probabilities for each unique context 28 | for context in unique_contexts: 29 | context_cleaned = re.sub(' +', ' ', context.replace('\n', ' ').strip()) 30 | 31 | print("\n------------------------------------") 32 | print(f"CONTEXT: {context_cleaned}") 33 | print("------------------------------------") 34 | 35 | generate_token_probabilities(model, tokenizer, context, top_k=top_k, max_depth=max_depth, additional_length=additional_length, prob_min=prob_min) 36 | 37 | def recursive_explore_and_generate(model, tokenizer, input_ids, prob_min, top_k, depth, max_depth, additional_length, device, path=[], pbar=None): 38 | # Initialize tqdm progress bar at the top level of recursion 39 | if pbar is None and depth == 0: 40 | # Initial total can be less than the maximum possible count 41 | initial_total = sum(math.pow(top_k, i) for i in range(depth + 1)) 42 | pbar = tqdm(total=initial_total, desc="Mapping tokens", leave=False) 43 | 44 | logits = model(input_ids).logits 45 | log_probabilities = torch.log_softmax(logits[0, -1], dim=0) 46 | top_k_tokens = torch.topk(log_probabilities, top_k).indices 47 | 48 | extended_token_probabilities = [] 49 | branches_explored = 0 50 | 51 | for token in top_k_tokens: 52 | token_prob = math.exp(log_probabilities[token].item()) 53 | 54 | # Update the tqdm progress bar 55 | pbar.update(1) 56 | 57 | # Only proceed if the token probability is greater than or equal to 10% 58 | if token_prob * 100 >= prob_min: 59 | branch_input_ids = torch.cat((input_ids, token.unsqueeze(0).unsqueeze(0)), dim=1) 60 | new_path = path + [(token.item(), token_prob)] 61 | branches_explored += 1 62 | 63 | if depth < max_depth: 64 | deeper_token_probs = recursive_explore_and_generate(model, tokenizer, branch_input_ids, prob_min, top_k, depth + 1, max_depth, additional_length, device, new_path, pbar) 65 | extended_token_probabilities.extend(deeper_token_probs) 66 | else: 67 | for _ in range(additional_length): 68 | logits = model(branch_input_ids).logits 69 | next_token_id = torch.argmax(logits[:, -1, :], dim=-1).unsqueeze(0) 70 | branch_input_ids = torch.cat((branch_input_ids, next_token_id), dim=1) 71 | 72 | sequence = tokenizer.decode(branch_input_ids[0], skip_special_tokens=True) 73 | sequence_cleaned = re.sub(' +', ' ', sequence.replace('\n', ' ').strip()) 74 | extended_token_probabilities.append((new_path, sequence_cleaned)) 75 | elif depth == max_depth - 1: 76 | pbar.update(1) 77 | 78 | # Adjust the total count of pbar based on actual branches explored 79 | if depth == 0 and branches_explored < top_k: 80 | pbar.total -= (top_k - branches_explored) * math.pow(top_k, max_depth - depth - 1) 81 | pbar.refresh() 82 | 83 | # Close the tqdm progress bar when the top-level recursion is finished 84 | if depth == 0: 85 | pbar.close() 86 | 87 | return extended_token_probabilities 88 | 89 | def colorize_probability(prob): 90 | """Colorize the probability percentage with a green-yellow-orange-red scale.""" 91 | if prob >= 80.0: 92 | color = Fore.LIGHTGREEN_EX 93 | elif prob >= 60.0: 94 | color = Fore.GREEN 95 | elif prob >= 40.0: 96 | color = Fore.LIGHTYELLOW_EX 97 | elif prob >= 20.0: 98 | color = Fore.YELLOW 99 | elif prob >= 10.0: 100 | color = Fore.LIGHTRED_EX 101 | else: 102 | color = Fore.RED 103 | return f'{color}{prob:.2f}%' 104 | 105 | def display_token_probabilities(token_probabilities, tokenizer): 106 | branches = 0 107 | 108 | for path, sequence in token_probabilities: 109 | branches += 1 110 | # Construct the probability path string with tokens 111 | prob_path_str = " > ".join([ 112 | f'{colorize_probability(token_prob * 100)}% [{tokenizer.decode([token], skip_special_tokens=True)}]' 113 | for token, token_prob in path 114 | ]) 115 | prob_path_str = prob_path_str.replace('\n', '\\n') 116 | print(f"Sequence: {sequence}\nTrace: {prob_path_str}{Fore.RESET}\n---") 117 | 118 | print(f"TOTAL BRANCHES: {branches}\n-----") 119 | 120 | def generate_token_probabilities(model, tokenizer, context="", prob_min=10, top_k=3, max_depth=5, additional_length=20, device='cuda'): 121 | input_ids = tokenizer.encode(context, return_tensors='pt').to(device) 122 | explored_and_extended_contexts = recursive_explore_and_generate(model, tokenizer, input_ids, prob_min, top_k, 0, max_depth, additional_length, device, []) 123 | display_token_probabilities(explored_and_extended_contexts, tokenizer) 124 | 125 | 126 | -------------------------------------------------------------------------------- /modules/merging.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | 4 | def merge_tensors(method: str, v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor: 5 | if method == "lerp": 6 | return merge_tensors_lerp(v0, v1, t) 7 | elif method == "slerp": 8 | return merge_tensors_slerp(v0, v1, t) 9 | elif method == "slice": 10 | return merge_tensors_slice(v0, v1, t) 11 | elif method == "cyclic": 12 | return merge_tensors_cyclic(v0, v1, t) 13 | elif method == "gradient": 14 | return merge_tensors_gradient(v0, v1, t) 15 | 16 | def merge_tensors_lerp(v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor: 17 | """Linear interpolation between two tensors.""" 18 | 19 | result = ((1 - t) * v0) + (t * v1) 20 | 21 | return result 22 | 23 | def merge_tensors_slerp(v0: torch.Tensor, v1: torch.Tensor, t: float, dot_threshold: float = 0.9995, eps: float = 1e-8) -> torch.Tensor: 24 | """Spherical linear interpolation between two tensors or linear interpolation if they are one-dimensional. 25 | Full credit to https://github.com/cg123/mergekit for the original code.""" 26 | 27 | # We LERP single dimensional tensors 28 | if v0.dim() == 1 and v1.dim() == 1: 29 | return merge_tensors_lerp(v0, v1, t) 30 | 31 | # Make copies of the original tensors to use for interpolation 32 | v0_copy = v0.clone() 33 | v1_copy = v1.clone() 34 | 35 | # Normalize the original tensors for angle computation 36 | v0 = safe_normalize(v0, eps) 37 | v1 = safe_normalize(v1, eps) 38 | 39 | # Compute the cosine of the angle between the normalized vectors. 40 | dot = (v0 * v1).sum() 41 | 42 | # If the inputs are too close, linearly interpolate using the original tensors. 43 | if abs(dot) > dot_threshold: 44 | return merge_tensors_lerp(v0_copy, v1_copy, t) 45 | 46 | # Calculate initial angle between v0 and v1 47 | theta_0 = torch.acos(dot) 48 | sin_theta_0 = torch.sin(theta_0) 49 | 50 | # Angle at timestep t 51 | theta_t = theta_0 * t 52 | sin_theta_t = torch.sin(theta_t) 53 | 54 | # Finish the slerp algorithm 55 | s0 = torch.sin(theta_0 - theta_t) / sin_theta_0 56 | s1 = sin_theta_t / sin_theta_0 57 | 58 | # Use the weights with the original tensors (not normalized) for the final result 59 | result = s0 * v0_copy + s1 * v1_copy 60 | 61 | return result 62 | 63 | # MODEL 1 > 10% blend > MODEL 2 64 | def merge_tensors_slice(v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor: 65 | # We're only working on the second dimension here 66 | if v0.dim() == 2: 67 | # Calculate the slice indices for each tensor 68 | slice_index_0 = int(v0.shape[1] * (1 - t)) 69 | slice_index_1 = v1.shape[1] - slice_index_0 70 | 71 | blend_slice_size = int(v0.shape[1] * 0.05) 72 | blend_slice_0 = v0.narrow(1, slice_index_0 - blend_slice_size, blend_slice_size * 2) 73 | blend_slice_1 = v1.narrow(1, slice_index_0 - blend_slice_size, blend_slice_size * 2) 74 | blended_slice = blend_slice_0 75 | 76 | # Apply gradient blending 77 | for i in range(blend_slice_size * 2): 78 | blend_ratio = i / (blend_slice_size * 2) 79 | blended_slice[:, i] = (blend_slice_1[:, i] * blend_ratio) + (blend_slice_0[:, i] * (1 - blend_ratio)) 80 | 81 | slice_index_0 = slice_index_0 - blend_slice_size 82 | slice_index_1 = slice_index_0 + blend_slice_size + blend_slice_size 83 | 84 | # Perform slicing 85 | slice_0 = v0.narrow(1, 0, slice_index_0) 86 | slice_1 = v1.narrow(1, slice_index_1, v1.shape[1] - slice_index_1) 87 | 88 | # Concatenate the slices 89 | result = torch.cat([slice_0, blended_slice, slice_1], dim=1) 90 | 91 | return result 92 | else: 93 | return v0 94 | 95 | # MODEL 1 > 10% blend > 10% of MODEL 2 > 10% blend > MODEL 1, with varying starting positions as defined by t 96 | def merge_tensors_cyclic(v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor: 97 | # We're only working on the second dimension here 98 | if v0.dim() == 2: 99 | blend_slice_size = int(v0.shape[1] * 0.05) # Blending zone is eventually multiplied by two due to overlap 100 | v1_slice_size = int(v0.shape[1] * 0.10) # 10% of Model 2, accounting for the 5% blend zone on both sides. So kinda 15%. 101 | 102 | slice_index_0 = int(v0.shape[1] * (1 - t)) - blend_slice_size # Model 1, first slice length 103 | 104 | # First MODEL 1 > MODEL 2 blend 105 | # ----------------------- 106 | blend_slice_0_0 = v0.narrow(1, slice_index_0, blend_slice_size * 2) 107 | blend_slice_0_1 = v1.narrow(1, slice_index_0, blend_slice_size * 2) 108 | blended_slice_0 = blend_slice_0_0 109 | 110 | # Apply gradient blending 111 | for i in range(blend_slice_size * 2): 112 | blend_ratio = i / (blend_slice_size * 2) 113 | blended_slice_0[:, i] = (blend_slice_0_0[:, i] * (1 - blend_ratio)) + (blend_slice_0_1[:, i] * blend_ratio) 114 | 115 | # Second MODEL 2 > MODEL 1 blend 116 | # ----------------------- 117 | blend_slice_1_0 = v0.narrow(1, slice_index_0 + (blend_slice_size * 2) + v1_slice_size, blend_slice_size * 2) 118 | blend_slice_1_1 = v1.narrow(1, slice_index_0 + (blend_slice_size * 2) + v1_slice_size, blend_slice_size * 2) 119 | blended_slice_1 = blend_slice_1_0 120 | 121 | # Apply gradient blending 122 | for i in range(blend_slice_size * 2): 123 | blend_ratio = i / (blend_slice_size * 2) 124 | blended_slice_1[:, i] = (blend_slice_1_1[:, i] * (1 - blend_ratio)) + (blend_slice_1_0[:, i] * blend_ratio) 125 | 126 | # Time to out main candidates into various pieces 127 | m1len_0 = slice_index_0 128 | m2start = slice_index_0 + (blend_slice_size * 2) 129 | m1start_1 = m2start + v1_slice_size + (blend_slice_size * 2) 130 | m2end_1 = v1.shape[1] - m1start_1 131 | 132 | # print(f"M1 0-{m1len_0} > B1 {m1len_0}-{m1len_0+(blend_slice_size * 2)} > M2 {m2start}-{m2start+v1_slice_size} > B2 {m2start+v1_slice_size}-{m1start_1} > M1 {m1start_1}-{v1.shape[1]}") 133 | 134 | slice_0_0 = v0.narrow(1, 0, m1len_0) # Model 1, first piece 135 | slice_1_0 = v1.narrow(1, m2start, v1_slice_size) # Model 2 slice 136 | slice_0_1 = v0.narrow(1, m1start_1, m2end_1) # Model 1, second piece 137 | 138 | # Concatenate the slices 139 | result = torch.cat([slice_0_0, blended_slice_0, slice_1_0, blended_slice_1, slice_0_1], dim=1) 140 | 141 | return result 142 | else: 143 | return v0 144 | 145 | # Model 1 > Model 2 > Model 1, with t defining the peak of the gradient along the tensor's width 146 | def merge_tensors_gradient(v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor: 147 | if v0.dim() == 2: 148 | total_length = v0.shape[1] 149 | peak = int(total_length * (1 - t)) 150 | 151 | # Create an index array 152 | indices = torch.arange(total_length).float() 153 | 154 | # Vectorized computation of blend ratios 155 | blend_ratios = torch.zeros_like(indices) 156 | blend_ratios[:peak] = (indices[:peak] / peak) * 0.9 # Scale to max 0.9 for v1 157 | blend_ratios[peak:] = torch.flip(indices[:total_length - peak], dims=[0]) / (total_length - peak) * 0.9 # Scale to max 0.9 for v1 158 | 159 | # Ensure that v0 still has influence 160 | v0_ratios = 1 - blend_ratios 161 | 162 | # Vectorized blending of the tensors 163 | result = (v1 * blend_ratios.unsqueeze(0)) + (v0 * v0_ratios.unsqueeze(0)) 164 | 165 | return result 166 | else: 167 | return v0 168 | 169 | def safe_normalize(tensor: torch.Tensor, eps: float): 170 | norm = tensor.norm() 171 | if norm > eps: 172 | return tensor / norm 173 | return tensor 174 | 175 | def merge_header_tensors(model1, model2, method, v0, v1, t) -> torch.Tensor: 176 | # TLDR - We reshape model 2's tensors to match model 1's 177 | model1bos = model1.config.bos_token_id 178 | model1eos = model1.config.eos_token_id 179 | model1size = v0.shape[0] 180 | 181 | model2bos = model2.config.bos_token_id 182 | model2eos = model2.config.eos_token_id 183 | model2size = v1.shape[0] 184 | 185 | # If model 2 has a smaller vocab, expand it 186 | if model1size > model2size: 187 | # Calculate the difference in size 188 | size_diff = model1size - model2size 189 | # Copy the additional entries from v0 to v1 190 | v1 = torch.cat([v1, v0[-size_diff:]], dim=0) 191 | 192 | # Swap special tokens if needed 193 | if model1bos != model2bos: 194 | v1[model1bos] = v1[model2bos] 195 | v1[model2bos] = v0[model1bos] 196 | if model1eos != model2eos: 197 | v1[model1eos] = v1[model2eos] 198 | v1[model2eos] = v0[model1eos] 199 | 200 | # If model 1 is smaller then 2, truncate 201 | # We do this after swapping tokens around 202 | if model1size < model2size: 203 | v1 = v1[:model1size] 204 | 205 | return merge_tensors_lerp(v0, v1, t) -------------------------------------------------------------------------------- /modules/models.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | import torch 4 | import transformers 5 | from transformers import AutoModelForCausalLM, AutoTokenizer 6 | from datetime import datetime 7 | 8 | class NoInit: 9 | def __enter__(self): 10 | def noop(*args, **kwargs): 11 | pass 12 | 13 | (k, u, n) = ( 14 | torch.nn.init.kaiming_uniform_, 15 | torch.nn.init.uniform_, 16 | torch.nn.init.normal_, 17 | ) 18 | torch.nn.init.kaiming_uniform_ = noop 19 | torch.nn.init.uniform_ = noop 20 | torch.nn.init.normal_ = noop 21 | 22 | transformers.modeling_utils._init_weights = False 23 | self.funcs = (k, u, n) 24 | 25 | def __exit__(self, *args): 26 | (k, u, n) = self.funcs 27 | ( 28 | torch.nn.init.kaiming_uniform_, 29 | torch.nn.init.uniform_, 30 | torch.nn.init.normal_, 31 | ) = ( 32 | k, 33 | u, 34 | n, 35 | ) 36 | transformers.modeling_utils._init_weights = True 37 | 38 | def load_model(model_path, device): 39 | with NoInit(): 40 | print("------------------------------------") 41 | print(f"{datetime.now().strftime('%H:%M:%S')} - Loading model ({model_path})...") 42 | 43 | if device == "cuda": 44 | tf_model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, device_map="auto") 45 | tf_model.half() 46 | else: 47 | tf_model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True) 48 | tf_model.half() 49 | tf_model = tf_model.to(device) 50 | 51 | tf_model.eval() 52 | print(f"{datetime.now().strftime('%H:%M:%S')} - Model loaded. Dtype: {tf_model.dtype}") 53 | print("------------------------------------") 54 | 55 | return tf_model 56 | 57 | def save_model(output_directory, source_directory, model): 58 | print(f"{datetime.now().strftime('%H:%M:%S')} - Saving model to {output_directory}...") 59 | model.save_pretrained(output_directory) 60 | 61 | # Check if additional files need to be copied 62 | if output_directory != source_directory: 63 | print(f"{datetime.now().strftime('%H:%M:%S')} - Copying tokenizer files to {output_directory}...") 64 | files_to_copy = ["added_tokens.json", "tokenizer.model", "special_tokens_map.json", 65 | "tokenizer_config.json", "vocab.json", "merges.txt", "tokenizer.json"] 66 | 67 | for filename in files_to_copy: 68 | src_path = os.path.join(source_directory, filename) 69 | dst_path = os.path.join(output_directory, filename) 70 | 71 | if os.path.exists(src_path): 72 | shutil.copy2(src_path, dst_path) 73 | print(f"Copied {filename}") 74 | else: 75 | print(f"Skipped {filename} (not found)") 76 | 77 | print(f"{datetime.now().strftime('%H:%M:%S')} - Model and tokenizer files saved successfully.") -------------------------------------------------------------------------------- /modules/probability.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import gc 3 | import math 4 | import torch 5 | 6 | from modules.utils import format_context 7 | 8 | # Global dictionary to store initial probabilities 9 | initial_phrase_probabilities = {} 10 | 11 | def convert_to_new_phrase_format(phrase_list): 12 | new_format_list = [] 13 | for entry in phrase_list: 14 | phrase = entry['phrase'] 15 | weight = entry.get('weight', 1) # Default weight to 1 if not present 16 | contexts = entry['contexts'] 17 | 18 | if isinstance(contexts[0], dict): # If already in the new format, copy as is 19 | new_format_list.append(entry) 20 | else: # Convert to the new format 21 | new_contexts = [{"context": context, "weight": weight} for context in contexts] 22 | new_format_list.append({"phrase": phrase, "contexts": new_contexts}) 23 | 24 | return new_format_list 25 | 26 | def auto_adjust_weights(model, tokenizer, bad_phrases, good_phrases, device): 27 | if device != "cuda": 28 | model_copy = copy.deepcopy(model).to('cuda') 29 | else: 30 | model_copy = model 31 | 32 | def adjust_phrase_weights(phrase_list): 33 | for entry in phrase_list: 34 | phrase = entry['phrase'] 35 | for context_entry in entry['contexts']: 36 | context = context_entry['context'] 37 | # Calculate unweighted joint probability 38 | joint_log_prob = calculate_joint_log_probability(model_copy, tokenizer, context, phrase) 39 | joint_prob = math.exp(joint_log_prob) 40 | 41 | # Adjust the weight: aiming for each phrase-context pair to have an equal contribution 42 | # Avoid division by zero; if joint_prob is 0, we can keep the weight unchanged 43 | if joint_prob > 0: 44 | context_entry['weight'] = 1.0 / joint_prob 45 | 46 | # Adjust weights for both bad and good phrases 47 | adjust_phrase_weights(bad_phrases) 48 | adjust_phrase_weights(good_phrases) 49 | 50 | if device != "cuda": 51 | del model_copy 52 | torch.cuda.empty_cache() 53 | 54 | return bad_phrases, good_phrases 55 | 56 | def calculate_joint_log_probability(model, tokenizer, context, phrase): 57 | sequence = context + phrase 58 | sequence_input_ids = tokenizer.encode(sequence, return_tensors="pt").to("cuda") 59 | 60 | with torch.no_grad(): 61 | outputs = model(sequence_input_ids) 62 | logits = outputs.logits 63 | 64 | phrase_tokens = tokenizer.encode(phrase, add_special_tokens=False) 65 | joint_log_prob = 0.0 66 | context_len = len(sequence_input_ids[0]) - len(phrase_tokens) 67 | 68 | for i, token_id in enumerate(phrase_tokens): 69 | word_log_prob = torch.log_softmax(logits[0, context_len + i - 1], dim=0)[token_id].item() 70 | joint_log_prob += word_log_prob 71 | 72 | return joint_log_prob 73 | 74 | def print_phrase_probabilities(model, tokenizer, bad_phrases, good_phrases, device): 75 | global initial_phrase_probabilities 76 | 77 | if device != "cuda": 78 | model_copy = copy.deepcopy(model).to('cuda') 79 | else: 80 | model_copy = model 81 | 82 | print("\n-----------------------------------------------------------------------------------------------------") 83 | print("| Type | Phrase | Context | Raw Prob* | Used Prob** | Change |") 84 | print("-----------------------------------------------------------------------------------------------------") 85 | 86 | # Initialize sums for good and bad phrases separately 87 | sums = { 88 | "BAD": {"real": 0, "weighted": 0, "change": 0}, 89 | "GOOD": {"real": 0, "weighted": 0, "change": 0} 90 | } 91 | 92 | for phrase_type, phrase_list in [("BAD", bad_phrases), ("GOOD", good_phrases)]: 93 | for entry in phrase_list: 94 | phrase = entry['phrase'] 95 | for context_entry in entry['contexts']: 96 | context = context_entry['context'] 97 | weight = context_entry['weight'] 98 | joint_log_prob = calculate_joint_log_probability(model_copy, tokenizer, context, phrase) 99 | joint_prob = math.exp(joint_log_prob) 100 | weighted_prob = joint_prob * weight 101 | 102 | # Update the sums 103 | sums[phrase_type]["real"] += joint_prob 104 | sums[phrase_type]["weighted"] += weighted_prob 105 | 106 | real_prob_str = f"{joint_prob * 100:.5f}%".ljust(12) 107 | 108 | if weighted_prob < 999999: prob_str = f"{weighted_prob * 100:.2f}%".ljust(12) 109 | else: prob_str = '###'.ljust(12) 110 | 111 | formatted_context = format_context(context.replace('\n',' '), 24) 112 | formatted_phrase = format_context(phrase.replace('\n',' '), 18) 113 | phrase_context_key = (phrase, context) 114 | 115 | if phrase_context_key not in initial_phrase_probabilities: 116 | initial_phrase_probabilities[phrase_context_key] = joint_prob 117 | print(f"| {phrase_type.ljust(4)} | {formatted_phrase} | {formatted_context} | {real_prob_str} | {prob_str} | {'N/A'.ljust(12)} |") 118 | else: 119 | initial_prob = initial_phrase_probabilities[phrase_context_key] 120 | change = ((joint_prob - initial_prob) * 100) * weight 121 | sums[phrase_type]["change"] += change 122 | 123 | if change < 999999: change_str = f"{change:+.2f}%".ljust(12) 124 | else: change_str = '###'.ljust(12) 125 | 126 | print(f"| {phrase_type.ljust(4)} | {formatted_phrase} | {formatted_context} | {real_prob_str} | {prob_str} | {change_str} |") 127 | 128 | # Calculate the net sums and print them 129 | net_real = sums["GOOD"]["real"] + sums["BAD"]["real"] 130 | net_weighted = sums["GOOD"]["weighted"] + sums["BAD"]["weighted"] 131 | net_change = sums["GOOD"]["change"] + sums["BAD"]["change"] 132 | 133 | net_real_str = f"{net_real * 100:.2f}%".ljust(12) 134 | 135 | if net_weighted < 999999: net_weighted_str = f"{net_weighted * 100:.2f}%".ljust(12) 136 | else: net_weighted_str = '###'.ljust(12) 137 | 138 | if net_change < 999999: net_change_str = f"{net_change:.2f}%".ljust(12) 139 | else: net_change_str = '###'.ljust(12) 140 | 141 | print("------------------------------------------------------------------------------------------------------") 142 | print(f"| {'Totals'.ljust(52)} | {net_real_str} | {net_weighted_str} | {net_change_str} |") 143 | print("------------------------------------------------------------------------------------------------------") 144 | print("* = Unweighted, raw probability - ** = Probability after weight adjustments\n") 145 | 146 | if device != "cuda": 147 | del model_copy 148 | torch.cuda.empty_cache() 149 | gc.collect() 150 | 151 | 152 | def calculate_word_probabilities(model, tokenizer, bad_phrases, good_phrases, device): 153 | if device != "cuda": 154 | model_copy = copy.deepcopy(model).to('cuda') 155 | else: 156 | model_copy = model 157 | 158 | phrase_probs = [] 159 | 160 | for phrase_list, sign in [(bad_phrases, 1), (good_phrases, -1)]: 161 | for entry in phrase_list: 162 | phrase = entry['phrase'] 163 | for context_entry in entry['contexts']: 164 | context = context_entry['context'] 165 | weight = context_entry['weight'] 166 | joint_log_prob = calculate_joint_log_probability(model_copy, tokenizer, context, phrase) 167 | joint_prob = math.exp(joint_log_prob) 168 | weighted_prob = joint_prob * weight * sign 169 | phrase_probs.append((phrase, context, weighted_prob)) 170 | 171 | if device != "cuda": 172 | del model_copy 173 | torch.cuda.empty_cache() 174 | gc.collect() 175 | 176 | return phrase_probs -------------------------------------------------------------------------------- /modules/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import yaml 3 | 4 | def load_config(config_path): 5 | with open(config_path, 'r') as file: 6 | return yaml.safe_load(file) 7 | 8 | class PrintAndStoreLogger: 9 | def __init__(self, original_stdout): 10 | self.contents = '' 11 | self.original_stdout = original_stdout 12 | 13 | def write(self, text): 14 | self.contents += text 15 | self.original_stdout.write(text) # Print to the console as well 16 | 17 | def flush(self): 18 | pass # This might be needed depending on the environment 19 | 20 | def print_ascii_art(file_path): 21 | try: 22 | with open(file_path, 'r') as file: 23 | ascii_art = file.read() 24 | print(ascii_art) 25 | except FileNotFoundError: 26 | print("ASCII art file not found.") 27 | 28 | def format_context(context, length=30): 29 | return (context[:length-2] + '..') if len(context) > length else context.ljust(length) -------------------------------------------------------------------------------- /monster-mapper.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import copy 3 | import gc 4 | import os 5 | import random 6 | import sys 7 | import torch 8 | import shutil 9 | import transformers 10 | import yaml 11 | 12 | from datetime import datetime 13 | from tqdm import tqdm 14 | from transformers import AutoModelForCausalLM, AutoTokenizer 15 | 16 | from modules.utils import print_ascii_art, format_context, load_config, PrintAndStoreLogger 17 | from modules.mapping import map_contexts 18 | from modules.models import load_model, NoInit 19 | from modules.probability import print_phrase_probabilities, convert_to_new_phrase_format 20 | 21 | def monster_mapper(config_path): 22 | # We save everything that gets printed to the screen 23 | original_stdout = sys.stdout 24 | logger = PrintAndStoreLogger(original_stdout) 25 | sys.stdout = logger # Redirect stdout to the logger instance 26 | 27 | config = load_config(config_path) 28 | 29 | if 'device' in config: device = config['device'] 30 | else: device = ['cpu'] 31 | 32 | model_path1 = config['directories']['model_path1'] 33 | 34 | # Phrase config 35 | if 'bad_phrases' in config: bad_phrases = config['bad_phrases'] 36 | else: bad_phrases = [] 37 | 38 | if 'good_phrases' in config: good_phrases = config['good_phrases'] 39 | else: good_phrases = [] 40 | 41 | # Seed config 42 | if 'random_seed' in config: random_seed = config['random_seed'] 43 | else: random_seed = 512 44 | 45 | # Mapper specific options 46 | if 'prob_min' in config['mapper']: prob_min = config['mapper']['prob_min'] 47 | else: prob_min = 10 48 | if 'top_k' in config['mapper']: top_k = config['mapper']['top_k'] 49 | else: top_k = 3 50 | if 'max_depth' in config['mapper']: max_depth = config['mapper']['max_depth'] 51 | else: max_depth = 10 52 | if 'additional_length' in config['mapper']: additional_length = config['mapper']['additional_length'] 53 | else: additional_length = 20 54 | 55 | # Actual start of script 56 | print_ascii_art("modules/logo.ascii") 57 | print(f"{datetime.now().strftime('%H:%M:%S')} - MONSTER CONTEXT MAPPER") 58 | print("------------------------------------") 59 | print(f"Device : {device}") 60 | print(f"Random seed : {random_seed}") 61 | print(f"Model to map : {model_path1}") 62 | print(f"Phrases loaded : {len(bad_phrases)+len(good_phrases)}") 63 | print("------------------------------------") 64 | print(f"Minimum branching prob : {prob_min}%") 65 | print(f"Top # per branch : {top_k}") 66 | print(f"Max branch depth : {max_depth}") 67 | print(f"Extra tokens generated : {additional_length}") 68 | 69 | with torch.no_grad(): 70 | if device == "cpu": torch.set_default_dtype(torch.float32) 71 | else: torch.set_default_dtype(torch.float16) 72 | 73 | torch.set_default_device(device) 74 | 75 | # Setting all the seeds 76 | torch.manual_seed(random_seed) 77 | random.seed(random_seed) 78 | torch.cuda.manual_seed(random_seed) 79 | torch.cuda.manual_seed_all(random_seed) 80 | torch.backends.cudnn.deterministic = True 81 | torch.backends.cudnn.benchmark = False 82 | 83 | # Testing the output model 84 | # model_path1 = output_directory 85 | 86 | # Load the base model + tokenizer 87 | model1 = load_model(model_path1, device) 88 | model1name = model_path1.split('/')[-1] 89 | header_chosen = [1.0, model1name] 90 | 91 | tokenizer = AutoTokenizer.from_pretrained(model_path1) 92 | tokenizer.padding_side = 'left' 93 | 94 | # Convert to new internal phrase format 95 | bad_phrases = convert_to_new_phrase_format(bad_phrases) 96 | good_phrases = convert_to_new_phrase_format(good_phrases) 97 | 98 | if device != "cuda": 99 | model1 = model1.to('cuda') 100 | 101 | # Mapping time! 102 | map_contexts(model1, tokenizer, bad_phrases, good_phrases, prob_min, top_k, max_depth, additional_length) 103 | 104 | # Save log 105 | sys.stdout = original_stdout # Restore the original stdout 106 | timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") 107 | 108 | # Create 'logs' subdirectory if it doesn't exist 109 | logs_dir = 'logs' 110 | if not os.path.exists(logs_dir): 111 | os.makedirs(logs_dir) 112 | 113 | # Define the log file path 114 | log_file_path = os.path.join(logs_dir, f"monster-mapper-{timestamp}.txt") 115 | 116 | # Write the log contents to the file in the 'logs' subdirectory 117 | with open(log_file_path, "w") as file: 118 | file.write(logger.contents) 119 | 120 | def main(): 121 | parser = argparse.ArgumentParser(description="Gryphe's Mythical Monster Mapper") 122 | parser.add_argument('--config', type=str, default='default.yaml', help='Path to the config YAML file') 123 | args = parser.parse_args() 124 | 125 | monster_mapper(args.config) 126 | 127 | if __name__ == "__main__": 128 | main() --------------------------------------------------------------------------------