Update app.py

app.py

@@ -1,52 +1,20 @@
-import os
-os.system("pip3 install torch==2.2.1 torchvision torchaudio   xformers --index-url https://download.pytorch.org/whl/cu121")
-import torch
-major_version, minor_version = torch.cuda.get_device_capability()
-# Must install separately since Colab has torch 2.2.1, which breaks packages
-os.system("pip install unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git")
-if major_version >= 8:
-    # Use this for new GPUs like Ampere, Hopper GPUs (RTX 30xx, RTX 40xx, A100, H100, L40)
-    os.system("pip install --no-deps packaging ninja einops flash-attn xformers trl peft \
-    accelerate bitsandbytes")
-else:
-    # Use this for older GPUs (V100, Tesla T4, RTX 20xx)
-    os.system("pip install --no-deps  trl peft accelerate bitsandbytes") 
-pass
-#os.system("git clone https://github.com/TimDettmers/bitsandbytes.git")
-#os.system("cd bitsandbytes/ && pip install -r requirements-dev.txt && cmake -DCOMPUTE_BACKEND=cuda -S . && make && pip install .")
-# Check if GPU is available
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
+import spaces  
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#print(f"Using device: {device}")
+device="cuda"
 model_name = "ruslanmv/Medical-Llama3-8B"
 device_map = 'auto'
-if device.type == "cuda":
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.float16,
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        quantization_config=bnb_config,
-        trust_remote_code=True,
-        use_cache=False,
-        device_map=device_map
-    )
-else:
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-
- # Load model directly
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
+model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
+@spaces.GPU  # Decorate the askme function with @spaces.GPU
 def askme(symptoms, question):
-    sys_message = '''
+    sys_message = '''\
     You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
     provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.
     '''
@@ -60,11 +28,11 @@ def askme(symptoms, question):
     return answer
 
 # Example usage
-symptoms = '''
+symptoms = '''\
 I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,
 increased sensitivity to cold, and dry, itchy skin.
 '''
-question = '''
+question = '''\
 Could these symptoms be related to hypothyroidism?
 If so, what steps should I take to get a proper diagnosis and discuss treatment options?
 '''