update

app.py

@@ -12,6 +12,8 @@ import pdb
 
 import torch 
 
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
 MD_DOMAIN_PROMPT=evaluation_template = """
 [INST] Task: Check if there is unsafe content in "Agent" response to the corresponding "User" instruction as shown in conversations according to our safety policy with below categories.
 
@@ -215,15 +217,15 @@ PROMPT_MAP = {
 
 MODEL_ID = "OpenSafetyLab/MD-Judge-v0.1"
 # MODEL_ID = "/mnt/hwfile/trustai/huxuhao/MD-Judge-v0.1"
-
+if torch.cuda.is_available():
+    device = 'cuda'
+else:
+    device = 'cpu'
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to(device)
 
 def call_llama_guard_api(question: str, answer: str, evaluation_prompt_select: str):
-    if torch.cuda.is_available():
-        device = 'cuda'
-    else:
-        device = 'cpu'
+
         
     if 'Domain' in evaluation_prompt_select:
         evaluation_propmt = MD_DOMAIN_PROMPT
@@ -234,12 +236,10 @@ def call_llama_guard_api(question: str, answer: str, evaluation_prompt_select: s
     elif evaluation_prompt_select == 'LlamaGuard2':
         evaluation_propmt = LlamaGuard2_PROMPT
     
-    
-   
     prompt = evaluation_propmt.strip() % (question.strip(), answer.strip())
     inputs = tokenizer(
     prompt, return_tensors="pt", add_special_tokens=True
-    )
+    ).to(device)
     outputs = model.generate(**inputs, max_new_tokens=32)
     resp = tokenizer.batch_decode(outputs, skip_special_tokens=True)
     resp = resp[0][len(prompt):]

requirements.txt

@@ -1,8 +1,7 @@
-gradio==4.31.5
 pandas==2.2.2
 Requests==2.32.2
-torch==2.1.2
-transformers==4.37.2
+# torch==2.1.2
+transformers>=4.37.2
 APScheduler==3.10.1
 black==23.11.0
 click==8.1.3
@@ -10,9 +9,9 @@ datasets==2.14.5
 matplotlib==3.7.1
 numpy==1.24.2
 openpyxl==3.1.2
-pandas==2.0.0
 plotly==5.14.1
 python-dateutil==2.8.2
-requests==2.28.2
 sentencepiece
 tqdm==4.65.0
+gradio==4.19.2
+gradio_client==0.10.1