OrangeEye commited on
Commit
654e004
·
1 Parent(s): 09a98a9
Files changed (2) hide show
  1. app.py +11 -11
  2. utils.py +1 -1
app.py CHANGED
@@ -35,7 +35,7 @@ generate_kwargs = dict(
35
  )
36
 
37
  # Load llama Guard
38
- llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID = load_llama_guard("meta-llama/Llama-Guard-3-1B")
39
 
40
  ## RAG MODEL
41
  RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert", n_gpu=0)
@@ -143,17 +143,17 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
143
 
144
  # @spaces.GPU(duration=60)
145
  def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'Trust-Align-Qwen2.5'):
146
- chat_round = [
147
- {"role": "user",
148
- "content": [
149
- {"type": "text",
150
- "text": message
151
- }
152
- ]
153
- }
154
- ]
155
  # llama guard check for it
156
- prompt_safety = moderate(chat_round, llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID)['generated_text']
157
  prompt_safety = "safe"
158
 
159
  if prompt_safety == "safe":
 
35
  )
36
 
37
  # Load llama Guard
38
+ # llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID = load_llama_guard("meta-llama/Llama-Guard-3-1B")
39
 
40
  ## RAG MODEL
41
  RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert", n_gpu=0)
 
143
 
144
  # @spaces.GPU(duration=60)
145
  def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'Trust-Align-Qwen2.5'):
146
+ # chat_round = [
147
+ # {"role": "user",
148
+ # "content": [
149
+ # {"type": "text",
150
+ # "text": message
151
+ # }
152
+ # ]
153
+ # }
154
+ # ]
155
  # llama guard check for it
156
+ # prompt_safety = moderate(chat_round, llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID)['generated_text']
157
  prompt_safety = "safe"
158
 
159
  if prompt_safety == "safe":
utils.py CHANGED
@@ -166,7 +166,7 @@ def load_llama_guard(model_id = "meta-llama/Llama-Guard-3-1B"):
166
  return llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID
167
 
168
 
169
- @spaces.GPU(duration=120)
170
  def moderate(chat, model, tokenizer, UNSAFE_TOKEN_ID):
171
 
172
  prompt = tokenizer.apply_chat_template(chat, return_tensors="pt", tokenize=False)
 
166
  return llama_guard, llama_guard_tokenizer, UNSAFE_TOKEN_ID
167
 
168
 
169
+ # @spaces.GPU(duration=120)
170
  def moderate(chat, model, tokenizer, UNSAFE_TOKEN_ID):
171
 
172
  prompt = tokenizer.apply_chat_template(chat, return_tensors="pt", tokenize=False)