wassname/Open-Assistant
1
0
Fork 0
mirror of https://github.com/wassname/Open-Assistant.git synced 2026-09-10 11:41:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

pre commits

Browse Source
This commit is contained in:
Sotirios Anagnostidis
2023-01-06 22:09:24 +01:00
parent fb3266b759
commit d3952354e2
7 changed files with 106 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/supervised_finetuning/models/__init__.py
+7 -5
View File
@@ -1,5 +1,6 @@
from transformers import AutoModelForCausalLM
from .gptj import get_model as get_gptj_model
# from .gptj import get_model as get_gptj_model
SUPPORTED_MODELS = ["galactica", "gpt-j"]
@@ -25,7 +26,8 @@ def freeze_top_n_layers(model, target_layers):
def get_specific_model(model_name, cache_dir, quantization):
if "gpt-j" in model_name.lower():
return get_gptj_model(model_name, cache_dir, quantization)
else:
return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)
return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)
# if "gpt-j" in model_name.lower():
# return get_gptj_model(model_name, cache_dir, quantization)
# else:
# return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)