Adjust GPU decorator

app.py

@@ -8,16 +8,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import os
 import spaces
 
-print("CUDA available:", torch.cuda.is_available())
-
-@spaces.GPU
-def claim_gpu():
-    # Dummy function to make Spaces detect GPU usage
-    pass
-
-claim_gpu()
-
-
 # Login automatically if HF_TOKEN is present
 hf_token = os.getenv("HF_TOKEN")
 if hf_token:
@@ -52,9 +42,10 @@ index.add(corpus_embeddings_np)
 
 # Generator (choose one: local HF model or OpenAI)
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
-model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", device_map="auto", torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", torch_dtype=torch.float16)
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150)
 
+@spaces.GPU
 def rag_pipeline(query):
     # Embed query
     query_embedding = embedder.encode([query], convert_to_tensor=True, device='cpu').numpy()
@@ -69,11 +60,7 @@ def rag_pipeline(query):
         print("-", repr(doc))
 
     # # Rerank
-    # rerank_pairs = [[str(query), str(doc)] for doc in retrieved_docs if isinstance(doc, str) and doc.strip()]
-    # if not rerank_pairs:
-    #     return "No valid documents found to rerank."
-    # scores = reranker.predict(rerank_pairs)
-
+    # rerank_pairs = [[str(query), str(doc)] for doc in retrieved_docs]
     # scores = reranker.predict(rerank_pairs)
     # reranked_docs = [doc for _, doc in sorted(zip(scores, retrieved_docs), reverse=True)]