Spaces:
Sleeping
Sleeping
File size: 4,704 Bytes
c83108e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import pandas as pd
import faiss
from sentence_transformers import SentenceTransformer
from transformers import T5ForConditionalGeneration, T5Tokenizer
import numpy as np
import gradio as gr
# --- 1. Load Models and Data (runs only once when the app starts) ---
print("Loading models and data... This may take a moment.")
# Load the dataset
df = pd.read_csv('syn5000.csv')
df.rename(columns={
'System / Subsystem Components': 'system',
'What is the item that you are focusing on?': 'item',
'What function does the item have?': 'function',
'What are you trying to achieve (Product Requirement)?': 'requirement',
'How could you get the requirements wrong (Failure Mode)?': 'failure_mode',
'Action Taken (Risk Mitigation)': 'mitigation'
}, inplace=True)
df['input_text'] = (
"System: " + df['system'] + "; " +
"Item: " + df['item'] + "; " +
"Requirement: " + df['requirement'] + "; " +
"Failure: " + df['failure_mode']
)
# Load the embedding model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
# Create and index embeddings using FAISS
corpus_embeddings = embedding_model.encode(df['input_text'].tolist())
embedding_dimension = corpus_embeddings.shape[1]
index = faiss.IndexFlatL2(embedding_dimension)
index.add(corpus_embeddings)
# Load the generator model and tokenizer
tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
generator_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
print("Models and data loaded successfully!")
# --- 2. The Core AI Logic ---
def retrieve_similar_examples(query_text, top_k=3):
query_embedding = embedding_model.encode([query_text])
distances, indices = index.search(query_embedding, top_k)
return df.iloc[indices[0]].to_dict('records')
def generate_mitigation_text(prompt):
inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
outputs = generator_model.generate(**inputs, max_length=128, num_beams=4, early_stopping=True)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# This is the main function that Gradio will call
def suggest_mitigation_from_ui(system, item, requirement, failure_mode):
"""
Takes individual text inputs from the UI and returns a suggested mitigation.
"""
query_text = (
f"System: {system}; "
f"Item: {item}; "
f"Requirement: {requirement}; "
f"Failure: {failure_mode}"
)
similar_examples = retrieve_similar_examples(query_text)
prompt = "You are an expert risk analysis engineer.\n\n"
prompt += "Based on the following similar past examples, write a specific risk mitigation action for the new failure described at the end.\n\n"
prompt += "--- EXAMPLES ---\n"
for ex in similar_examples:
prompt += f"Failure Description: {ex['input_text']}\n"
prompt += f"Mitigation Action: {ex['mitigation']}\n---\n"
prompt += "\n--- NEW FAILURE ---\n"
prompt += f"Failure Description: {query_text}\n"
prompt += "Mitigation Action:"
generated_text = generate_mitigation_text(prompt)
# We can also return the examples it used, for transparency
retrieved_info = "--- Retrieved Similar Examples ---\n"
for i, ex in enumerate(similar_examples):
retrieved_info += f"{i+1}. {ex['input_text'][:150]}...\n"
return generated_text, retrieved_info
# --- 3. Create the Gradio Web Interface ---
with gr.Blocks() as demo:
gr.Markdown("# AI Risk Mitigation Assistant")
gr.Markdown("Enter the details of a potential failure to get an AI-generated mitigation suggestion based on historical data.")
with gr.Row():
with gr.Column():
system_input = gr.Textbox(label="System / Subsystem")
item_input = gr.Textbox(label="Item in Focus")
requirement_input = gr.Textbox(label="Product Requirement")
failure_mode_input = gr.Textbox(label="Failure Mode")
submit_btn = gr.Button("Suggest Mitigation", variant="primary")
with gr.Column():
output_mitigation = gr.Textbox(label="✅ AI-Generated Mitigation Suggestion", lines=5)
output_examples = gr.Textbox(label="Retrieved Examples", lines=5)
submit_btn.click(
fn=suggest_mitigation_from_ui,
inputs=[system_input, item_input, requirement_input, failure_mode_input],
outputs=[output_mitigation, output_examples]
)
# This launches the app. On Hugging Face, it will be served automatically.
if __name__ == "__main__":
demo.launch() |