Spaces:
Running
Running
import os | |
import gradio as gr | |
import logging | |
from executor import risk_identifier, mitigations | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
class UI: | |
def __init__(self): | |
"""Load any static assets """ | |
self.load_css() | |
def header_block(self): | |
"""Title/description""" | |
with open("static/text/header.md", "r") as f: | |
content = f.read() | |
gr.Markdown(content) | |
gr.Markdown("---") | |
gr.Markdown("<br>") | |
def risk_extraction(self): | |
with open("static/text/sec1.md") as f: | |
content = f.read() | |
with gr.Row(): | |
gr.Markdown(content) | |
with gr.Row(): | |
with gr.Column(variant="compact", scale=1): | |
self.usecase = gr.TextArea( | |
label="Intent", | |
interactive=True, | |
info="Describe the intent of the application, or choose from one of the examples below." | |
) | |
self.taxonomy = gr.Dropdown( | |
choices=["ibm-risk-atlas", "mit-ai-risk-repository", "nist-ai-rmf", "ibm-granite-guardian"], | |
multiselect=False, | |
value="ibm-risk-atlas", | |
label="Choose a risk taxonomy.", | |
info="The risk taxonomy defines a wide range of risks, their classifications, and potential mitigations.", | |
interactive=True, | |
) | |
self.model_name_or_path = gr.Dropdown( | |
#choices=['codellama/codellama-34b-instruct-hf', 'google/flan-t5-xl', 'google/flan-t5-xxl', 'google/flan-ul2', 'ibm/granite-13b-instruct-v2', 'ibm/granite-20b-code-instruct', 'ibm/granite-20b-multilingual', 'ibm/granite-3-2-8b-instruct-preview-rc', 'ibm/granite-3-2b-instruct', 'ibm/granite-3-8b-instruct', 'ibm/granite-34b-code-instruct', 'ibm/granite-3b-code-instruct', 'ibm/granite-8b-code-instruct', 'ibm/granite-guardian-3-2b', 'ibm/granite-guardian-3-8b', 'meta-llama/llama-2-13b-chat', 'meta-llama/llama-3-1-70b-instruct', 'meta-llama/llama-3-1-8b-instruct', 'meta-llama/llama-3-2-11b-vision-instruct', 'meta-llama/llama-3-2-1b-instruct', 'meta-llama/llama-3-2-3b-instruct', 'meta-llama/llama-3-2-90b-vision-instruct', 'meta-llama/llama-3-3-70b-instruct', 'meta-llama/llama-3-405b-instruct', 'meta-llama/llama-guard-3-11b-vision', 'mistralai/mistral-large', 'mistralai/mixtral-8x7b-instruct-v01'], | |
choices=["ibm/granite-20b-code-instruct"], | |
value="ibm/granite-20b-code-instruct", | |
multiselect=False, | |
label="Choose language model to use", | |
info="Language model used to assess risks (This is not the model being assessed).", | |
interactive=True | |
) | |
examples = gr.Examples([["In a medical chatbot, generative AI can be employed to create a triage system that assesses patients' symptoms and provides immediate, contextually relevant advice based on their medical history and current condition. The chatbot can analyze the patient's input, identify potential medical issues, and offer tailored recommendations or insights to the patient or healthcare provider. This can help streamline the triage process, ensuring that patients receive the appropriate level of care and attention, and ultimately improving patient outcomes.", "ibm-risk-atlas"], | |
["Building a customer support agent that automatically triages common problems with services.", "ibm-risk-atlas"]], | |
[self.usecase, self.taxonomy], | |
label='Example use cases', example_labels=["Medical chatbot", "Customer service agent"] | |
) | |
self.risk_execute = gr.Button("Submit") | |
with gr.Column(scale=2): | |
self.assessment_sec = gr.Markdown() | |
self.assessed_risks = gr.Dataset(label=None, visible=False) | |
self.assessed_risk_definition = gr.Markdown() | |
gr.Markdown( | |
"""<h2> Related Risks </h2> | |
Select a potential risk above to check for related risks. | |
""" | |
) | |
rrtb = gr.Markdown() | |
self.relatedrisks = gr.Dataset(components=[rrtb], label=None, visible=False) | |
gr.Markdown( | |
"""<h2> Mitigations </h2> | |
Select a potential risk to determine possible mitigations. """ | |
) | |
self.mitigations_text = gr.Markdown() | |
self.mitigations = gr.DataFrame(label=None, visible=False) | |
self.download = gr.DownloadButton("Download JSON", visible=False) | |
gr.Markdown("---") | |
gr.Markdown("<br>") | |
def load_css(self): | |
with open("static/style.css", "r") as file: | |
self.css = file.read() | |
def layout(self): | |
"""Assemble the overall layout""" | |
with gr.Blocks(theme=gr.themes.Default()) as demo: # type: ignore | |
self.header_block() | |
self.risks = gr.State() | |
# Risk assessment based on user intents | |
self.risk_extraction() | |
# Register event listeners | |
self.risk_execute.click( | |
fn=risk_identifier, | |
inputs=[ | |
self.usecase, | |
self.model_name_or_path, | |
self.taxonomy | |
], | |
outputs=[self.assessment_sec, self.risks, self.assessed_risks, self.download], | |
api_name="risk_identifier" | |
) | |
self.assessed_risks.select( | |
fn=mitigations, | |
inputs=[self.assessed_risks, self.taxonomy], | |
# NOTETOSELF: Intent based risk is stored in self.risk (if needed) | |
outputs=[self.assessed_risk_definition, self.relatedrisks, self.mitigations, self.mitigations_text] | |
) | |
return demo | |
def run(self): | |
self.ui = self.layout() | |
self.ui.queue().launch(allowed_paths=["static/"], ssr_mode=False) | |
if __name__ == "__main__": | |
ui = UI() | |
ui.run() | |