PhantHive commited on
Commit
49282d8
·
verified ·
1 Parent(s): 4cf26ce

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -14,10 +14,11 @@ model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
14
 
15
  # Load the Lora model
16
  model = PeftModel.from_pretrained(model, model_id)
 
17
 
18
  def greet(text):
19
  with torch.no_grad(): # Disable gradient calculation for inference
20
- batch = tokenizer(text, return_tensors='pt').to(device) # Move tensors to device
21
  with torch.cuda.amp.autocast(): # Enable mixed-precision if available
22
  output_tokens = model.generate(**batch, max_new_tokens=15)
23
  return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
 
14
 
15
  # Load the Lora model
16
  model = PeftModel.from_pretrained(model, model_id)
17
+ model.to(device)
18
 
19
  def greet(text):
20
  with torch.no_grad(): # Disable gradient calculation for inference
21
+ batch = tokenizer(f'"{text}" ->:', return_tensors='pt').to(device) # Move tensors to device
22
  with torch.cuda.amp.autocast(): # Enable mixed-precision if available
23
  output_tokens = model.generate(**batch, max_new_tokens=15)
24
  return tokenizer.decode(output_tokens[0], skip_special_tokens=True)