34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
from Context import Context
|
|
from ModelLoader import ModelLoader
|
|
from transformers import AutoTokenizer
|
|
from peft import PeftModel
|
|
|
|
class Merger:
|
|
def mergeAndStore(self, context: Context):
|
|
if not context.mergeFull:
|
|
return
|
|
if not context.storeAdapter:
|
|
raise Exception("Configuration problem: mergeFull requires storeAdapter")
|
|
|
|
print("Reloading model")
|
|
baseModel = ModelLoader().load(context.locBaseModel, False, True, context)
|
|
peftModel = PeftModel.from_pretrained(baseModel, context.locAdapter)
|
|
|
|
print("Merging")
|
|
mergedModel = peftModel.merge_and_unload()
|
|
|
|
print("Dequantize")
|
|
try:
|
|
mergedModel = mergedModel.dequantize()
|
|
except:
|
|
pass
|
|
|
|
print("Storing model")
|
|
mergedModel.save_pretrained(save_directory = context.locFull)
|
|
|
|
print("Storing tokenizer")
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
context.locBaseModel,
|
|
trust_remote_code = True)
|
|
tokenizer.save_pretrained(context.locFull)
|
|
|