Spaces:
Runtime error
Runtime error
File size: 4,146 Bytes
4710d1d |
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 |
import streamlit as st
import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration, BitsAndBytesConfig, CLIPProcessor, CLIPModel
from peft import PeftModel
from PIL import Image
st.set_page_config(page_title="Multimodal Risk Engine", page_icon="🛡️", layout="wide")
# --- LOAD MODELS (Smart Caching) ---
@st.cache_resource
def load_models():
print("🔄 Loading Models...")
device = "cuda" if torch.cuda.is_available() else "cpu"
# A. Stage 1 (CLIP)
clip_id = "openai/clip-vit-base-patch32"
clip_model = CLIPModel.from_pretrained(clip_id).to(device)
clip_processor = CLIPProcessor.from_pretrained(clip_id)
# B. Stage 2 (LLaVA)
model_id = "llava-hf/llava-1.5-7b-hf"
# CPU vs GPU Loading Logic
if device == "cuda":
bnb_config = BitsAndBytesConfig(
load_in_4bit=True, bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16
)
base_model = LlavaForConditionalGeneration.from_pretrained(
model_id, quantization_config=bnb_config, torch_dtype=torch.float16, device_map="auto"
)
else:
# Fallback for Free CPU Tier (Might be slow/crash but allows build verification)
base_model = LlavaForConditionalGeneration.from_pretrained(model_id)
adapter_id = "oke39/llava-v1.5-7b-hateful-memes-lora"
model = PeftModel.from_pretrained(base_model, adapter_id)
llava_processor = AutoProcessor.from_pretrained(model_id)
return clip_model, clip_processor, model, llava_processor
try:
clip_model, clip_processor, llava_model, llava_processor = load_models()
st.toast("✅ System Ready", icon="🚀")
except Exception as e:
st.error(f"Hardware Error: {e}")
st.stop()
# --- PIPELINE ---
def stage_1_glance(image):
HATE_ANCHOR = ["hate speech, offensive content, racism, dangerous propaganda"]
inputs = clip_processor(text=HATE_ANCHOR, images=image, return_tensors="pt", padding=True).to(clip_model.device)
with torch.no_grad():
outputs = clip_model(**inputs)
probs = outputs.logits_per_image.softmax(dim=1)
return float(probs[0][0])
def stage_2_judge(image, text_context):
device = "cuda" if torch.cuda.is_available() else "cpu"
prompt = f"USER: <image>\nAnalyze this meme text: '{text_context}'. Is it hateful? Return JSON.\nASSISTANT:"
inputs = llava_processor(text=prompt, images=image, return_tensors="pt").to(device)
with torch.inference_mode():
output = llava_model.generate(**inputs, max_new_tokens=200)
response = llava_processor.decode(output[0], skip_special_tokens=True)
return response.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in response else response
# --- UI ---
st.title("🛡️ Multimodal Content Risk Engine")
st.markdown("### The 'Benign Confounder' Solver")
col1, col2 = st.columns([1, 1])
with col1:
uploaded_file = st.file_uploader("Upload Meme", type=["jpg", "png", "jpeg"])
text_input = st.text_input("Extracted Text", placeholder="Type visible text...")
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
st.image(image, caption="User Upload", use_container_width=True)
with col2:
if uploaded_file and st.button("Analyze Risk", type="primary"):
with st.status("Running Pipeline...", expanded=True) as status:
st.write("👁️ **Stage 1: The 'Glance' (CLIP)**")
risk_score = stage_1_glance(image)
st.progress(min(risk_score, 1.0))
if risk_score < 0.22:
status.update(label="✅ Auto-Approved", state="complete")
st.success("Verdict: BENIGN")
st.write(f"Risk Score: `{risk_score:.4f}`")
else:
st.warning(f"⚠️ Risk Detected ({risk_score:.4f})! Escalating...")
st.write("⚖️ **Stage 2: The 'Judge' (LLaVA)**")
verdict = stage_2_judge(image, text_input if text_input else "")
status.update(label="✅ Done", state="complete")
st.json(verdict)
|