import json
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import numpy as np

categories = ["Very Safe", "Safe", "Slightly Risky", "Moderate", "Risky", "Very Risky", "Highly Dangerous"]


# Load pre-trained DistilBERT model and tokenizer
# Load pre-trained DistilBERT model and tokenizer
model_name = "distilbert-base-uncased"
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=len(categories))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Function to classify a single description
def classify_description(model, tokenizer, device, description):
    inputs = tokenizer(description, padding=True, truncation=True, return_tensors="pt")
    input_ids = inputs["input_ids"].to(device)

    with torch.no_grad():
        outputs = model(input_ids)
    logits = outputs.logits

    # Simulating a more granular classification (for illustration)
    predicted_label = np.random.choice(len(categories))  # Replace with actual model logic
    return categories[predicted_label]

# Streamlit app to upload JSON and display results
st.title("HTML Content Classifier")

uploaded_file = st.file_uploader("Upload JSON file", type=["json"])
if uploaded_file is not None:
    # Read the file
    json_data = json.load(uploaded_file)

    # Display results
    for item in json_data:
        classification = classify_description(item["description"])
        st.write(f"Title: {item['title']}")
        st.write(f"Link: {item['link']}")
        st.write(f"Description: {item['description']}")
        st.write(f"Classification: {classification}")
        st.write("------")