Update app.py

app.py

@@ -6,6 +6,7 @@ from transformers import pipeline
 from huggingface_hub import login
 from streamlit.components.v1 import html
 import pandas as pd
+import torch
 
 # Retrieve the token from environment variables
 hf_token = os.environ.get("HF_TOKEN")
@@ -58,9 +59,9 @@ st.write("This model will score your reviews in your CSV file and generate a rep
 @st.cache_resource
 def load_models():
     # Load the scoring model via pipeline.
-    score_pipe = pipeline("text-classification", model="mixedbread-ai/mxbai-rerank-base-v1")
+    score_pipe = pipeline("text-classification", model="mixedbread-ai/mxbai-rerank-base-v1", device=0)
     # Load the Gemma text generation pipeline.
-    gemma_pipe = pipeline("text-generation", model="google/gemma-3-1b-it")
+    gemma_pipe = pipeline("text-generation", model="google/gemma-3-1b-it", device=0)
     return score_pipe, gemma_pipe
 
 score_pipe, gemma_pipe = load_models()