deepspeed-chat: support any model in chatbot (#744)

Currently, chatbot assumes OPTForCausalLM model. Modify it to use the required model from the checkpoint. Change-Id: I04cbc28f87c7be4fc89a3fac39a3e5634b151b32 Signed-off-by: Moshe Island <[email protected]> Co-authored-by: Moshe Island <[email protected]>
microsoft · Oct 2, 2023 · d8f3f73 · d8f3f73
1 parent 9b3d898
commit d8f3f73
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/applications/DeepSpeed-Chat/inference/chatbot.py b/applications/DeepSpeed-Chat/inference/chatbot.py
@@ -10,7 +10,7 @@
 import os
 import json
 from transformers import pipeline, set_seed
-from transformers import AutoConfig, OPTForCausalLM, AutoTokenizer
+from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
 
 
 def parse_args():
@@ -43,9 +43,10 @@ def get_generator(path):
     tokenizer.pad_token = tokenizer.eos_token
 
     model_config = AutoConfig.from_pretrained(path)
-    model = OPTForCausalLM.from_pretrained(path,
-                                           from_tf=bool(".ckpt" in path),
-                                           config=model_config).half()
+    model_class = AutoModelForCausalLM.from_config(model_config)
+    model = model_class.from_pretrained(path,
+                                        from_tf=bool(".ckpt" in path),
+                                        config=model_config).half()
 
     model.config.end_token_id = tokenizer.eos_token_id
     model.config.pad_token_id = model.config.eos_token_id