artificialguybr commited on
Commit
ddc101d
1 Parent(s): 61c1d8a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import requests
3
+ import os
4
+ import json
5
+
6
+ api_key = os.getenv('API_KEY')
7
+
8
+ def call_llama_guard_api(content, assistant_response):
9
+ invoke_url = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/b34280ac-24e4-4081-bfaa-501e9ee16b6f"
10
+ headers = {
11
+ "Authorization": f"Bearer {api_key}",
12
+ "Accept": "application/json",
13
+ }
14
+ payload = {
15
+ "messages": [
16
+ {"content": content, "role": "user"},
17
+ {"content": assistant_response, "role": "assistant"}
18
+ ]
19
+ }
20
+
21
+ session = requests.Session()
22
+ response = session.post(invoke_url, headers=headers, json=payload)
23
+
24
+ while response.status_code == 202:
25
+ request_id = response.headers.get("NVCF-REQID")
26
+ fetch_url = f"https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/{request_id}"
27
+ response = session.get(fetch_url, headers=headers)
28
+
29
+ response.raise_for_status()
30
+ response_body = response.json()
31
+ print(response_body)
32
+ return response_body
33
+
34
+ content_input = gr.Textbox(lines=2, placeholder="Enter your content here...", label="User Content")
35
+ assistant_response_input = gr.Textbox(lines=2, placeholder="Enter assistant's response here...", label="Assistant Response")
36
+
37
+ iface = gr.Interface(fn=call_llama_guard_api,
38
+ inputs=[content_input, assistant_response_input],
39
+ outputs="text",
40
+ title="Llama Guard Safety Classifier",
41
+ description="Classify the safety of LLM prompts and responses using Llama Guard"
42
+ )
43
+
44
+ iface.launch()