import os
from dotenv import load_dotenv
import gradio as gr
from transformers import pipeline
from huggingface_hub import login

load_dotenv()

hf_token = os.getenv('HUGGING_FACE_TOKEN')

if hf_token:
    login(token=hf_token)
else:
    raise ValueError("HUGGING_FACE_TOKEN is not set in the environment variables")

def classify_text(text):
    pipe = pipeline("text-classification", model="meta-llama/Prompt-Guard-86M")
    result = pipe(text)[0]
    return {result['label']: result['score']}

demo = gr.Interface(
    fn=classify_text,
    inputs="text",
    outputs="label",
    title="Prompt-Guard-86M Text Classification"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)