Update app.py

app.py

@@ -337,7 +337,7 @@ def optimize_query(
     vector_store_type: str,  # Added to match your signature
     search_type: str,  # Added to match your signature
     top_k: int = 3,
-    use_gpu: bool = False
+    use_gpu: bool = True
 ) -> str:
     """
     CPU-optimized version of query expansion using a small language model.
@@ -439,11 +439,11 @@ def optimize_query(
         if device == "cuda":
             torch.cuda.empty_cache()
             
-        return expanded_query.strip()
+        return [Document(page_content=expanded_query.strip())]
         
     except Exception as e:
         print(f"Query optimization failed: {str(e)}")
-        return query  # Return original query if optimization fails
+        return [Document(page_content=query)]  # Return original query if optimization fails