Kanpredict
/

gptj-6b-8bits

Text Generation

Inference Endpoints

Model card Files Files and versions Community

OssamaLafhel commited on Apr 5, 2023

Commit

4b15582

·

1 Parent(s): 859f302

Update handler.py

Files changed (1) hide show

handler.py +2 -2

handler.py CHANGED Viewed

@@ -159,8 +159,8 @@ transformers.models.gptj.modeling_gptj.GPTJBlock = GPTJBlock  # monkey-patch GPT
 # -----------------------------------------> API <---------------------------------------
 tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
-model = GPTJForCausalLM.from_pretrained("Kanpredict/gptj-6b-8bits", low_cpu_mem_usage=True)
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
 class EndpointHandler:

 # -----------------------------------------> API <---------------------------------------
 tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
+model = GPTJForCausalLM.from_pretrained("", low_cpu_mem_usage=True)
+device = 0 if torch.cuda.is_available() else -1
 class EndpointHandler: