Kowsher
/

TokenTrails

Text Generation

text-generation-inference

Model card Files Files and versions

Kowsher commited on Jul 18, 2023

Commit

13ad266

·

1 Parent(s): 1db61ef

Update handler.py

Files changed (1) hide show

handler.py +2 -1

handler.py CHANGED Viewed

@@ -2,8 +2,9 @@ import torch
 from typing import Any, Dict
 from transformers import AutoModelForCausalLM, AutoTokenizer
 class EndpointHandler:
     def __init__(self, path=""):
         # load model and tokenizer from path

 from typing import Any, Dict
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.models.auto import modeling_auto
+modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES['falcon'] = 'FalconForCausalLM'
 class EndpointHandler:
     def __init__(self, path=""):
         # load model and tokenizer from path