28 lines
1015 B
Python
28 lines
1015 B
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 original model on CPU")
|
|
baseModel = ModelLoader().load(context.locBaseModel, True, False, context)
|
|
peftModel = PeftModel.from_pretrained(baseModel, context.locAdapter)
|
|
|
|
print("Merging")
|
|
mergedModel = peftModel.merge_and_unload()
|
|
|
|
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)
|
|
|