Kowsher commited on
Commit
13ad266
1 Parent(s): 1db61ef

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +2 -1
handler.py CHANGED
@@ -2,8 +2,9 @@ import torch
2
 
3
  from typing import Any, Dict
4
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
5
 
6
-
7
  class EndpointHandler:
8
  def __init__(self, path=""):
9
  # load model and tokenizer from path
 
2
 
3
  from typing import Any, Dict
4
  from transformers import AutoModelForCausalLM, AutoTokenizer
5
+ from transformers.models.auto import modeling_auto
6
 
7
+ modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES['falcon'] = 'FalconForCausalLM'
8
  class EndpointHandler:
9
  def __init__(self, path=""):
10
  # load model and tokenizer from path