Use CUDA

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.utils import logging
 
@@ -8,10 +9,18 @@ logger = logging.get_logger("transformers")
 
 # Load model directly from your Hugging Face repository
 def load_model():    
-
     tokenizer = AutoTokenizer.from_pretrained("omi-health/sum-small", trust_remote_code=False)
     model = AutoModelForCausalLM.from_pretrained("omi-health/sum-small", trust_remote_code=False)
     
+    # Move model to GPU if available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    
+    print(f"Using device: {device}")
+    if device == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+        print(f"Memory allocated: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
+    
     return model, tokenizer
 
 def generate_soap_note(doctor_patient_conversation):
@@ -34,6 +43,8 @@ Please generate a structured SOAP (Subjective, Objective, Assessment, Plan) note
         max_length=tokenizer.model_max_length
     )
     
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
     generate_ids = model.generate(
         inputs.input_ids,
         attention_mask=inputs.attention_mask,  # Explicitly pass attention mask

requirements.txt

@@ -1,4 +1,4 @@
---extra-index-url https://download.pytorch.org/whl/cu113
+--extra-index-url https://download.pytorch.org/whl/cu118
 torch
 transformers>=4.36.0
 gradio>=3.50.0