Update app.py

app.py

@@ -14,24 +14,14 @@ from kvpress import (
 )
 import spaces
 
-# Initialize GPU Zero
-zero = torch.Tensor([0]).cuda()
-print(zero.device)  # Ensure the tensor is on GPU
-
 @spaces.GPU
 def process_request(url, question, press_type, compression_ratio):
     try:
-        print(zero.device)  # Confirm the GPU usage
-
         # Fetch Wikipedia content
         content = requests.get(url).content
         soup = BeautifulSoup(content, "html.parser")
         context = "".join([p.text for p in soup.find_all("p")]) + "\n\n"
 
-        # Calculate tokens
-        tokens = pipe.tokenizer.encode(context, return_tensors="pt").to(device)
-        num_tokens = tokens.size(1)
-
         # Initialize the press
         press_class = press_map.get(press_type)
         if not press_class:
@@ -49,7 +39,7 @@ def process_request(url, question, press_type, compression_ratio):
 # Load pipeline
 device = "cuda:0"
 ckpt = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-attn_implementation = "flash_attention_2"
+attn_implementation = "sdpa"
 pipe = pipeline(
     "kv-press-text-generation",
     model=ckpt,