Update app.py

app.py

@@ -12,11 +12,11 @@ bnb_config = BitsAndBytesConfig(
     )
 
 
-tokenizer = AutoTokenizer.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True).to(device)
 # Load model in this way if use GPU
-# model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True, quantization_config=bnb_config)
+model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True, quantization_config=bnb_config).to(device)
 
-model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True)
+# model = AutoModelForCausalLM.from_pretrained("truongghieu/deci-finetuned", trust_remote_code=True)
 # Move the model to the GPU if available
 
 generation_config = GenerationConfig(
@@ -34,7 +34,7 @@ generation_config = GenerationConfig(
 # Define a function that takes a text input and generates a text output
 def generate_text(text):
     input_text = f'###Human: \"{text}\"'
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
     output_ids = model.generate(input_ids, generation_config=generation_config)
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return output_text