Update app.py

app.py

@@ -1,24 +1,26 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
-from threading import Thread
 import os
 
-import torch
+# Check CUDA availability
 print("CUDA available:", torch.cuda.is_available())
 print("CUDA version:", torch.version.cuda)
 
-# Cargar el token de Hugging Face desde los secretos
+# Load the Hugging Face token from secrets
 token = os.environ.get("HF_TOKEN")
 
 model_id = "google/shieldgemma-2b"
 
-# use quantization to lower GPU usage
+# Use quantization to lower GPU usage
 bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_id,token=token)
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.bfloat16,