Update app.py

app.py

@@ -4,86 +4,132 @@ import json
 import logging
 import traceback
 import os
+import torch
 import numpy as np
 from datetime import datetime
-from transformers import pipeline, set_seed
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from sentence_transformers import SentenceTransformer, util
+from diffusers import StableDiffusionPipeline
+import librosa
+import soundfile as sf
+import tempfile
 
-# Import our components
+# ARF components
 from agentic_reliability_framework.runtime.engine import EnhancedReliabilityEngine
 from hallucination_detective import HallucinationDetectiveAgent
 from memory_drift_diagnostician import MemoryDriftDiagnosticianAgent
+from image_detector import ImageQualityDetector
+from audio_detector import AudioQualityDetector
 from ai_event import AIEvent
 from ai_risk_engine import AIRiskEngine
 from nli_detector import NLIDetector
+from retrieval import SimpleRetriever
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Initialize infrastructure engine (optional)
+# ----------------------------------------------------------------------
+# Infrastructure engine (optional)
+# ----------------------------------------------------------------------
 try:
     logger.info("Initializing EnhancedReliabilityEngine...")
     engine = EnhancedReliabilityEngine()
-    logger.info("Engine initialized successfully.")
 except Exception as e:
-    logger.error(f"Failed to initialize engine: {e}\n{traceback.format_exc()}")
+    logger.error(f"Engine init failed: {e}")
     engine = None
 
-# Load generative model (small autoregressive)
+# ----------------------------------------------------------------------
+# Generative model for text (DialoGPT-small)
+# ----------------------------------------------------------------------
 gen_model_name = "microsoft/DialoGPT-small"
+tokenizer = AutoTokenizer.from_pretrained(gen_model_name)
+model = AutoModelForCausalLM.from_pretrained(gen_model_name)
+logger.info(f"Generator {gen_model_name} loaded.")
+
+# ----------------------------------------------------------------------
+# NLI detector
+# ----------------------------------------------------------------------
+nli_detector = NLIDetector()
+
+# ----------------------------------------------------------------------
+# Sentence‑Transformer retriever
+# ----------------------------------------------------------------------
+retriever = SimpleRetriever()
+logger.info("Retriever loaded.")
+
+# ----------------------------------------------------------------------
+# Image generation (tiny model for demo)
+# ----------------------------------------------------------------------
 try:
-    generator = pipeline('text-generation', model=gen_model_name, device=0 if torch.cuda.is_available() else -1)
-    logger.info(f"Generator {gen_model_name} loaded.")
+    image_pipe = StableDiffusionPipeline.from_pretrained(
+        "hf-internal-testing/tiny-stable-diffusion-torch"
+    )
+    if not torch.cuda.is_available():
+        image_pipe.to("cpu")
+    logger.info("Image pipeline loaded.")
 except Exception as e:
-    logger.error(f"Failed to load generator: {e}")
-    generator = None
+    logger.error(f"Image pipeline failed: {e}")
+    image_pipe = None
 
-# Load NLI detector
-nli_detector = NLIDetector()
+# ----------------------------------------------------------------------
+# Audio transcription (Whisper tiny)
+# ----------------------------------------------------------------------
+try:
+    audio_pipe = pipeline(
+        "automatic-speech-recognition",
+        model="openai/whisper-tiny.en",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    logger.info("Audio pipeline loaded.")
+except Exception as e:
+    logger.error(f"Audio pipeline failed: {e}")
+    audio_pipe = None
 
+# ----------------------------------------------------------------------
 # AI agents
+# ----------------------------------------------------------------------
 hallucination_detective = HallucinationDetectiveAgent(nli_detector=nli_detector)
 memory_drift_diagnostician = MemoryDriftDiagnosticianAgent()
+image_quality_detector = ImageQualityDetector()
+audio_quality_detector = AudioQualityDetector()
 
-# AI risk engine
+# ----------------------------------------------------------------------
+# Bayesian risk engine
+# ----------------------------------------------------------------------
 ai_risk_engine = AIRiskEngine()
 
-# In‑memory storage for last event to attach feedback
-last_ai_event = None
-last_ai_category = None
-
-async def generate_response(prompt: str, max_length: int = 100) -> tuple:
-    """Generate response using the small autoregressive model."""
-    if generator is None:
-        return "[Model not loaded]", 0.0, "Model loading failed"
-    try:
-        loop = asyncio.get_event_loop()
-        # We need to compute confidence; text-generation pipeline returns text but not logits.
-        # For simplicity, we'll set confidence based on a heuristic (e.g., generation length?).
-        # Alternatively, use a model that returns probabilities.
-        # Let's use a simple placeholder: confidence = 0.8 if generation succeeds.
-        # In practice, we'd need to access logits.
-        result = await loop.run_in_executor(
-            None,
-            lambda: generator(prompt, max_new_tokens=max_length, return_full_text=False)
+# ----------------------------------------------------------------------
+# Generation helper with log probabilities
+# ----------------------------------------------------------------------
+def generate_with_logprobs(prompt, max_new_tokens=100):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            return_dict_in_generate=True,
+            output_scores=True
         )
-        response = result[0]['generated_text']
-        # Placeholder confidence
-        confidence = 0.8
-        return response, confidence, ""
-    except Exception as e:
-        logger.error(f"Generation error: {e}")
-        return "", 0.0, str(e)
+    scores = outputs.scores
+    log_probs = [torch.log_softmax(score, dim=-1) for score in scores]
+    generated_ids = outputs.sequences[0][inputs['input_ids'].shape[1]:]
+    token_log_probs = []
+    for i, lp in enumerate(log_probs):
+        token_id = generated_ids[i]
+        token_log_probs.append(lp[0, token_id].item())
+    avg_log_prob = sum(token_log_probs) / len(token_log_probs) if token_log_probs else 0.0
+    generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    return generated_text, avg_log_prob
 
-async def analyze_ai(task_type, prompt):
-    global last_ai_event, last_ai_category
+# ----------------------------------------------------------------------
+# Task handlers
+# ----------------------------------------------------------------------
+async def handle_text(task_type, prompt):
     try:
-        # Generate response
-        response, confidence, error = await generate_response(prompt)
-        if error:
-            return json.dumps({"error": error}, indent=2)
-
-        # Create AIEvent
+        response, avg_log_prob = generate_with_logprobs(prompt)
+        # Get retrieval score
+        retrieval_score = retriever.get_similarity(prompt)
+        # Create event
         event = AIEvent(
             timestamp=datetime.utcnow(),
             component="ai",
@@ -99,91 +145,173 @@ async def analyze_ai(task_type, prompt):
             prompt=prompt,
             response=response,
             response_length=len(response),
-            confidence=confidence,
+            confidence=float(np.exp(avg_log_prob)),  # convert to probability scale
             perplexity=None,
-            retrieval_scores=None,
+            retrieval_scores=[retrieval_score],
             user_feedback=None,
             latency_ms=0
         )
-        last_ai_event = event
-        last_ai_category = task_type
-
-        # Run agents
+        # Analyze
         hallu_result = await hallucination_detective.analyze(event)
         drift_result = await memory_drift_diagnostician.analyze(event)
-
-        # Get current risk metrics
         risk_metrics = ai_risk_engine.risk_score(task_type)
-
-        result = {
+        return {
             "response": response,
-            "confidence": confidence,
+            "avg_log_prob": avg_log_prob,
+            "confidence": event.confidence,
+            "retrieval_score": retrieval_score,
             "hallucination_detection": hallu_result,
             "memory_drift_detection": drift_result,
             "risk_metrics": risk_metrics
         }
-        return json.dumps(result, indent=2)
     except Exception as e:
-        logger.error(f"AI analysis error: {e}\n{traceback.format_exc()}")
-        return json.dumps({"error": str(e), "traceback": traceback.format_exc()}, indent=2)
+        logger.error(f"Text task error: {e}")
+        return {"error": str(e)}
+
+async def handle_image(prompt):
+    if image_pipe is None:
+        return {"error": "Image model not loaded"}
+    try:
+        import time
+        start = time.time()
+        image = image_pipe(prompt, num_inference_steps=2).images[0]  # tiny steps for speed
+        gen_time = time.time() - start
+        # Mock retrieval score (you could use CLIP similarity)
+        retrieval_score = retriever.get_similarity(prompt)
+        event = AIEvent(
+            timestamp=datetime.utcnow(),
+            component="image",
+            service_mesh="ai",
+            latency_p99=0,
+            error_rate=0.0,
+            throughput=1,
+            cpu_util=None,
+            memory_util=None,
+            action_category="image",
+            model_name="tiny-sd",
+            model_version="latest",
+            prompt=prompt,
+            response="",  # image not text
+            response_length=0,
+            confidence=1.0 / (gen_time + 1),  # heuristic
+            perplexity=None,
+            retrieval_scores=[retrieval_score, gen_time],
+            user_feedback=None,
+            latency_ms=gen_time * 1000
+        )
+        quality_result = await image_quality_detector.analyze(event)
+        return {
+            "image": image,
+            "generation_time": gen_time,
+            "retrieval_score": retrieval_score,
+            "quality_detection": quality_result
+        }
+    except Exception as e:
+        logger.error(f"Image task error: {e}")
+        return {"error": str(e)}
 
+async def handle_audio(audio_file):
+    if audio_pipe is None:
+        return {"error": "Audio model not loaded"}
+    try:
+        # Load audio (Gradio provides file path)
+        audio, sr = librosa.load(audio_file, sr=16000)
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            sf.write(tmp.name, audio, sr)
+            result = audio_pipe(tmp.name, return_timestamps=False)
+        text = result["text"]
+        # Whisper does not output log probs easily; we'll use a placeholder
+        avg_log_prob = -2.0  # placeholder
+        event = AIEvent(
+            timestamp=datetime.utcnow(),
+            component="audio",
+            service_mesh="ai",
+            latency_p99=0,
+            error_rate=0.0,
+            throughput=1,
+            cpu_util=None,
+            memory_util=None,
+            action_category="audio",
+            model_name="whisper-tiny.en",
+            model_version="latest",
+            prompt="",  # audio file path
+            response=text,
+            response_length=len(text),
+            confidence=float(np.exp(avg_log_prob)),
+            perplexity=None,
+            retrieval_scores=[avg_log_prob],
+            user_feedback=None,
+            latency_ms=0
+        )
+        quality_result = await audio_quality_detector.analyze(event)
+        return {
+            "transcription": text,
+            "avg_log_prob": avg_log_prob,
+            "confidence": event.confidence,
+            "quality_detection": quality_result
+        }
+    except Exception as e:
+        logger.error(f"Audio task error: {e}")
+        return {"error": str(e)}
+
+# ----------------------------------------------------------------------
+# Feedback handling
+# ----------------------------------------------------------------------
+last_event_category = None
 def feedback(thumbs_up: bool):
-    """Handle user feedback to update Beta priors."""
-    global last_ai_category, last_ai_event
-    if last_ai_category is None:
+    global last_event_category
+    if last_event_category is None:
         return "No previous analysis to rate."
-    ai_risk_engine.update_outcome(last_ai_category, success=thumbs_up)
-    # Optionally, also update the event with feedback
-    if last_ai_event:
-        last_ai_event.user_feedback = thumbs_up
-    return f"Feedback recorded: {'👍' if thumbs_up else '👎'} for {last_ai_category}."
+    ai_risk_engine.update_outcome(last_event_category, success=thumbs_up)
+    return f"Feedback recorded: {'👍' if thumbs_up else '👎'} for {last_event_category}."
 
-# Build the Gradio interface
+# ----------------------------------------------------------------------
+# Gradio UI
+# ----------------------------------------------------------------------
 with gr.Blocks(title="ARF v4 – AI Reliability Lab", theme="soft") as demo:
-    gr.Markdown("# 🧠 ARF v4 – AI Reliability Lab\n**Detect hallucinations and drift in generative AI**")
+    gr.Markdown("# 🧠 ARF v4 – AI Reliability Lab\n**Detect hallucinations, drift, and failures across text, image, and audio**")
     
-    with gr.Row():
-        with gr.Column():
-            task_type = gr.Dropdown(
-                choices=["chat", "code", "summary"],
-                value="chat",
-                label="Task Type"
-            )
-            prompt = gr.Textbox(
-                label="Prompt",
-                value="What is the capital of France?",
-                lines=3
-            )
-            analyze_btn = gr.Button("Analyze", variant="primary")
-        with gr.Column():
-            output = gr.JSON(label="Analysis Result")
+    with gr.Tabs():
+        with gr.TabItem("Text Generation"):
+            text_task = gr.Dropdown(["chat", "code", "summary"], value="chat", label="Task")
+            text_prompt = gr.Textbox(label="Prompt", value="What is the capital of France?")
+            text_btn = gr.Button("Generate")
+            text_output = gr.JSON(label="Analysis")
+        
+        with gr.TabItem("Image Generation"):
+            img_prompt = gr.Textbox(label="Prompt", value="A cat wearing a hat")
+            img_btn = gr.Button("Generate")
+            img_output = gr.Image(label="Generated Image")
+            img_json = gr.JSON(label="Analysis")
+        
+        with gr.TabItem("Audio Transcription"):
+            audio_input = gr.Audio(type="filepath", label="Upload audio file")
+            audio_btn = gr.Button("Transcribe")
+            audio_output = gr.JSON(label="Analysis")
     
     with gr.Row():
-        feedback_btn_up = gr.Button("👍 Correct")
-        feedback_btn_down = gr.Button("👎 Incorrect")
+        feedback_up = gr.Button("👍 Correct")
+        feedback_down = gr.Button("👎 Incorrect")
         feedback_msg = gr.Textbox(label="Feedback", interactive=False)
     
-    analyze_btn.click(
-        fn=analyze_ai,
-        inputs=[task_type, prompt],
-        outputs=output
+    # Wire up events
+    text_btn.click(
+        fn=lambda task, p: asyncio.run(handle_text(task, p)),
+        inputs=[text_task, text_prompt],
+        outputs=text_output
     )
-    feedback_btn_up.click(
-        fn=lambda: feedback(True),
-        outputs=feedback_msg
+    img_btn.click(
+        fn=lambda p: asyncio.run(handle_image(p)),
+        inputs=img_prompt,
+        outputs=[img_output, img_json]
     )
-    feedback_btn_down.click(
-        fn=lambda: feedback(False),
-        outputs=feedback_msg
+    audio_btn.click(
+        fn=lambda f: asyncio.run(handle_audio(f)),
+        inputs=audio_input,
+        outputs=audio_output
     )
-    
-    gr.Markdown("""
-    ---
-    - **Model**: `microsoft/DialoGPT-small` (autoregressive, 117M params)
-    - **NLI Detector**: `typeform/distilroberta-base-mnli` (82M params)
-    - **Risk engine**: Beta conjugate priors per task category
-    - **Feedback** updates the posterior distribution
-    """)
+    feedback_up.click(fn=lambda: feedback(True), outputs=feedback_msg)
+    feedback_down.click(fn=lambda: feedback(False), outputs=feedback_msg)
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)