thfname commited on
Commit
2b6957b
1 Parent(s): 04d4939

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -5,12 +5,12 @@ from transformers import AutoModelForCausalLM
5
  from transformers import AutoTokenizer
6
 
7
  model_name = "thfname/test-gemma2-2b"
8
- model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to("cuda")
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
 
11
 
12
  def gt (tx):
13
- input_ids = tokenizer(tx, return_tensors="pt").to("cuda")
14
  outputs = model.generate(**input_ids, max_new_tokens=32)
15
  return tokenizer.decode(outputs[0])
16
 
 
5
  from transformers import AutoTokenizer
6
 
7
  model_name = "thfname/test-gemma2-2b"
8
+ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
 
11
 
12
  def gt (tx):
13
+ input_ids = tokenizer(tx, return_tensors="pt")
14
  outputs = model.generate(**input_ids, max_new_tokens=32)
15
  return tokenizer.decode(outputs[0])
16