update

app.py

@@ -38,7 +38,7 @@ generate_kwargs = dict(
 llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID = load_llama_guard("meta-llama/Llama-Guard-3-1B")
 
 ## RAG MODEL
-RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
+RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert", n_gpu=1)
 
 try:
   gr.Info("Setting up retriever, please wait...")

requirements.txt

@@ -12,4 +12,5 @@ arxiv
 ragatouille
 hf_transfer
 colorlog
-accelerate==1.1.1
+accelerate==1.1.1
+faiss-gpu

utils.py

@@ -188,7 +188,7 @@ def moderate(chat, model, tokenizer, UNSAFE_TOKEN_ID):
     ######
     # Get generated text
     ######
-    logger.info(outputs)
+    # logger.info(outputs)
     # Number of tokens that correspond to the input prompt
     input_length = inputs.input_ids.shape[1]
     # Ignore the tokens from the input to get the tokens generated by the model