Update app.py
Browse files
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)
|