K00B404 commited on
Commit
ceda046
β€’
1 Parent(s): 6d26b7f

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +2 -2
model.py CHANGED
@@ -4,7 +4,7 @@ from typing import Iterator
4
  import torch
5
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
6
 
7
- model_id = 'codellama/CodeLlama-13b-Instruct-hf'
8
 
9
  if torch.cuda.is_available():
10
  config = AutoConfig.from_pretrained(model_id)
@@ -22,7 +22,7 @@ else:
22
  torch.set_default_tensor_type(torch.FloatTensor)
23
  model = AutoModelForCausalLM.from_pretrained(
24
  model_id,
25
- torch_dtype=torch.float32,
26
  device_map='cpu')
27
 
28
  tokenizer = AutoTokenizer.from_pretrained(model_id)
 
4
  import torch
5
  from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
6
 
7
+ model_id = 'TheBloke/CodeLlama-7B-GGML'
8
 
9
  if torch.cuda.is_available():
10
  config = AutoConfig.from_pretrained(model_id)
 
22
  torch.set_default_tensor_type(torch.FloatTensor)
23
  model = AutoModelForCausalLM.from_pretrained(
24
  model_id,
25
+ torch_dtype=torch.float16,
26
  device_map='cpu')
27
 
28
  tokenizer = AutoTokenizer.from_pretrained(model_id)