Link to WML + UI improvements

.gitignore

@@ -1 +1,4 @@
-.DS_Store
+.DS_Store
+.env
+scratch.py
+__pycache__

README.md

@@ -1,10 +1,10 @@
 ---
 title: Risk Atlas Nexus
-emoji: 🏆
+emoji: 🛡️
 colorFrom: purple
 colorTo: green
 sdk: gradio
-sdk_version: 5.16.2
+sdk_version: 5.18.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -30,7 +30,7 @@ class UI:
         with gr.Row():
             gr.Markdown(content)
 
-        with gr.Row(equal_height=True):
+        with gr.Row():
             with gr.Column(variant="compact", scale=1):
                 self.usecase = gr.TextArea(
                     label="Intent",
@@ -60,34 +60,24 @@ class UI:
                                        )
                 self.risk_execute = gr.Button("Submit")
         
-            with gr.Column(variant="panel", scale=2):
+            with gr.Column(scale=2):
             
-                gr.Markdown(
-                    """<h2 style='color: black;'> Potential Risks </h2> """
-                )
-                self.assessed_risks = gr.Dataset(label=None)
-            # gr.DataFrame(
-            #     headers=["name", "description", "concern"],
-            #     datatype=["str", "str", "str"], 
-            #     label="Risks",
-            #     visible=True, 
-            #     wrap=True,
-            #     show_copy_button=True,
-            #     show_search="filter",
-            #     type='pandas')
+                self.assessment_sec = gr.Markdown()
+                self.assessed_risks = gr.Dataset(label=None, visible=False)
 
                 gr.Markdown(
-                """<h2 style='color: black;'> Related Risks </h2> 
+                """<h2> Related Risks </h2> 
                 Select a potential risk above to check for related risks.
                 """
                 )
-                self.relatedrisks = gr.Dataset(label=None, visible=False)
+                rrtb = gr.Markdown()
+                self.relatedrisks = gr.Dataset(components=[rrtb], label=None, visible=False)
 
                 gr.Markdown(
-                """<h2 style='color: black;'> Mitigations </h2> 
+                """<h2> Mitigations </h2> 
                 Select a potential risk to determine possible mitigations. """
                 )
-                self.mitigations = gr.Dataset(label=None, visible=False)
+                self.mitigations = gr.DataFrame(label=None, visible=False)
             
         gr.Markdown("---")
         gr.Markdown("<br>")
@@ -100,7 +90,7 @@ class UI:
     def layout(self):
         """Assemble the overall layout""" 
 
-        with gr.Blocks(theme=gr.themes.Default()) as demo:
+        with gr.Blocks(theme=gr.themes.Default()) as demo: # type: ignore
 
             self.header_block()
             self.risks = gr.State()
@@ -116,7 +106,7 @@ class UI:
                     self.model_name_or_path,
                     self.taxonomy
                 ],
-                outputs=[self.risks, self.assessed_risks],
+                outputs=[self.assessment_sec, self.risks, self.assessed_risks],
                 api_name="risk_identifier"
             )
             self.assessed_risks.select(

executor.py

@@ -4,6 +4,7 @@ load_dotenv(override=True)
 
 import re
 import os
+import pandas as pd
 import json
 from typing import List, Dict, Any
 import pandas as pd
@@ -17,7 +18,7 @@ from risk_atlas_nexus.library import RiskAtlasNexus
 from functools import lru_cache
 
 # Load the taxonomies
-ran = RiskAtlasNexus()
+ran = RiskAtlasNexus() # type: ignore
 
 
 @lru_cache
@@ -25,53 +26,42 @@ def risk_identifier(usecase: str,
                     model_name_or_path: str = "ibm/granite-20b-code-instruct", 
                     taxonomy: str = "ibm-risk-atlas"): # -> List[Dict[str, Any]]: #pd.DataFrame:
 
-    # inference_engine = WMLInferenceEngine(
-    #     model_name_or_path= model_name_or_path,
-    #     credentials={
-    #         "api_key": os.environ["WML_API_KEY"],
-    #         "api_url": os.environ["WML_API_URL"],
-    #         "project_id": os.environ["WML_PROJECT_ID"],
-    #     },
-    #     parameters=WMLInferenceEngineParams(
-    #         max_new_tokens=100, decoding_method="greedy", repetition_penalty=1
-    #     ),  # type: ignore
-    # )
-
-    # risks = ran.identify_risks_from_usecase(
-    #     usecase=usecase,
-    #     inference_engine=inference_engine,
-    #     taxonomy=taxonomy,
-    # )
-    risks = ' ["Harmful code generation", "Hallucination", "Harmful output", "Toxic output", "Spreading toxicity", "Spreading disinformation", "Nonconsensual use", "Non-disclosure", "Data contamination", "Data acquisition restrictions", "Data usage rights restrictions", "Confidential data in prompt", "Confidential information in data", "Personal information in prompt", "Personal information in data", "IP information in prompt",'
-
-
-    if isinstance(risks, str):
-        # Translate LLM output to Risk catalog
-        out = []
-        try:
-            risks = json.loads(risks)
-        except json.JSONDecodeError:
-            
-            # Fallback to regex - will skip any partial categories
-            risks = re.findall(r'"(.*?)"', risks)
-
-        for risk in risks:
-            k = list(filter(lambda r: r.name == risk, ran._ontology.risks)) # type: ignore
-            out+= ([i.model_dump() for i in k])
-
-    elif isinstance(risks, List):
-        # FIXME: assumes that the output is structured - not sure if that's correct.
-        out = risks
-    
-    #out_df = pd.DataFrame(out)
+    inference_engine = WMLInferenceEngine(
+        model_name_or_path= model_name_or_path,
+        credentials={
+            "api_key": os.environ["WML_API_KEY"],
+            "api_url": os.environ["WML_API_URL"],
+            "project_id": os.environ["WML_PROJECT_ID"],
+        },
+        parameters=WMLInferenceEngineParams(
+            max_new_tokens=150, decoding_method="greedy", repetition_penalty=1
+        ),  # type: ignore
+    )
+
+    risks = ran.identify_risks_from_usecases( # type: ignore
+        usecases=[usecase],
+        inference_engine=inference_engine,
+        taxonomy=taxonomy,
+    )[0]
+
+    sample_labels = [r.name if r else r.id for r in risks]
+
+    out_sec = gr.Markdown("""<h2> Potential Risks </h2> """)
         
     #return out_df
-    return gr.State(out), gr.Dataset(samples=[i['id'] for i in out], 
-                                     sample_labels=[i['name'] for i in out], samples_per_page=50, visible=True, label="Estimated by an LLM.")
+    return out_sec, gr.State(risks), gr.Dataset(samples=[r.id for r in risks], 
+                                     sample_labels=sample_labels, 
+                                     samples_per_page=50, visible=True, label="Estimated by an LLM.")
     
 
 @lru_cache
-def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Dataset, gr.Dataset]:
+def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Dataset, gr.DataFrame]:
+    """
+    For a specific risk (riskid), returns
+    (a) related risks - as a dataset
+    (b) mitigations
+
+    """
     
     related_risk_ids = ran.get_related_risk_ids_by_risk_id(riskid)
     action_ids = []
@@ -96,19 +86,21 @@ def mitigations(riskid: str, taxonomy: str) -> tuple[gr.Dataset, gr.Dataset]:
         samples = None
         sample_labels = None
     else:
-        label = "Related risks"
+        label = f"Risks from other taxonomies related to {riskid}"
         samples = related_risk_ids
-        sample_labels = [i.name for i in ran.get_related_risks_by_risk_id(riskid)]
+        sample_labels = [i.name for i in ran.get_related_risks_by_risk_id(riskid)] #type: ignore
 
     if not action_ids:
         alabel = "No mitigations found."
         asamples = None
         asample_labels = None
+        mitdf = pd.DataFrame()
     else:
-        alabel = ""
+        alabel = f"Mitigation actions related to risk {riskid}."
         asamples = action_ids
-        asample_labels = [ran.get_action_by_id(i).description for i in action_ids]
-
+        asample_labels = [ran.get_action_by_id(i).description for i in asamples] # type: ignore
+        asample_name = [ran.get_action_by_id(i).name for i in asamples] #type: ignore
+        mitdf = pd.DataFrame({"Mitigation": asample_name, "Description": asample_labels})
 
     return (gr.Dataset(samples=samples, label=label, sample_labels=sample_labels, visible=True),
-            gr.Dataset(samples=asamples, label=alabel, sample_labels=asample_labels, visible=True))
+            gr.DataFrame(mitdf, wrap=True, show_copy_button=True, show_search="search", label=alabel, visible=True)) 

requirements.txt

@@ -1,4 +1,5 @@
-gradio==5.16.2
+gradio==5.18.0
 linkml==1.8.6
 linkml_runtime==1.8.3
-risk-atlas-nexus @ git+https://github.com/rahulnair23/risk-atlas-nexus
+ibm_watsonx_ai==1.2.8
+risk-atlas-nexus [wml] @ git+https://github.com/IBM/risk-atlas-nexus

static/text/header.md

@@ -1,7 +1,7 @@
 
 
 <img src="/gradio_api/file=static/images/Risk_Atlas_Nexus_vector.png" width=300 style="display:block;margin:0 auto;">
-<h1 style='text-align: center; color: black;'> Risk Atlas Nexus </h1>
+<h1 style='text-align: center;'> Risk Atlas Nexus </h1>
 
 
 Risk Atlas Nexus aims to provide tooling to help bring together disparate resources related to governance of foundation models. We support a community driven approach to curating and cataloguing resources such as datasets, benchmarks and mitigations. Our goal is to turn abstract risk definitions into actionable workflows that streamline AI governance processes. By connecting fragmented resources, Risk Atlas Nexus seeks to fill a critical gap in AI governance, enabling stakeholders to build more robust, transparent, and accountable systems. Risk Atlas Nexus builds on the [IBM AI Risk Atlas](https://www.ibm.com/docs/en/watsonx/saas?topic=ai-risk-atlas) making this educational resource a nexus of governance assets and tooling. An AI System's Knowledge Graph is used to provide a unified structure that links and contextualize the very heterogeneous domain data.

static/text/sec1.md

@@ -1,3 +1,3 @@
-<h1 style='text-align: left; color: black;'> Risk assessment based on intent </h1>
+<h1 style='text-align: left;'> Risk assessment based on intent </h1>
 
 Assess potential risks for a use case description based on risk taxonomies.