Update app.py

app.py

@@ -1,15 +1,20 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM , GenerationConfig
-
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 
+import torch
 
+# Check if a GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 tokenizer = AutoTokenizer.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
+model.to(device)
+
 generation_config = GenerationConfig(
     penalty_alpha=0.6,
-    do_sample = True,
+    do_sample=True,
     top_k=5,
     temperature=0.5,
     repetition_penalty=1.2,
@@ -17,11 +22,10 @@ generation_config = GenerationConfig(
     pad_token_id=tokenizer.eos_token_id
 )
 
-
 # Define a function that takes a text input and generates a text output
 def generate_text(text):
     input_text = text
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)  # Move input to the GPU
     output_ids = model.generate(input_ids, generation_config=generation_config)
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return output_text