import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

@st.cache_resource
def load_model():
    import os
    # Set cache dir to writable path
    os.environ["TRANSFORMERS_CACHE"] = "/app/hf-cache"
    os.makedirs("/app/hf-cache", exist_ok=True)

    model_name = "tiiuae/falcon-rw-1b"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, output_attentions=True, trust_remote_code=True)
    model.eval()
    return tokenizer, model



def get_attention(prompt, tokenizer, model):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    return outputs.attentions, tokenizer.convert_ids_to_tokens(inputs.input_ids[0])

def main():
    st.title("🔐 Basic PII Attention Auditor")

    tokenizer, model = load_model()

    prompt = st.text_area("Enter your prompt:", "My name is Alice and my SSN is 123-45-6789.")
    pii_keywords = ["ssn", "alice", "123", "email", "contact"]

    if st.button("Analyze"):
        attentions, tokens = get_attention(prompt, tokenizer, model)

        pii_indices = [i for i, t in enumerate(tokens) if any(k in t.lower() for k in pii_keywords)]
        st.write("Detected PII tokens:", [tokens[i] for i in pii_indices])

        st.success("Base pipeline working. Ready for advanced visualization in future upgrade.")

if __name__ == "__main__":
    main()