Reizend
/

bhai_insights_v3

Text Generation

text-generation-inference

Model card Files Files and versions Community

Reizendretail commited on Nov 6, 2023

Commit

3c18151

•

1 Parent(s): de2af79

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -5,7 +5,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model = AutoModelForCausalLM.from_pretrained(path)
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.tokenizer.use_default_system_prompt = False

 class EndpointHandler():
     def __init__(self, path=""):
+        self.model = AutoModelForCausalLM.from_pretrained(path,device_map="cuda:0", load_in_4bit=True)
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.tokenizer.use_default_system_prompt = False