import streamlit as st
import torch
from transformers import (GPT2Tokenizer, GPT2ForSequenceClassification)

labels_ids = {0: "Human Generated", 1: "AI Generated"}

model = GPT2ForSequenceClassification.from_pretrained(pretrained_model_name_or_path='ErnestBeckham/gpt-2-finetuned-ai-content')
tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model_name_or_path='ErnestBeckham/gpt2-tokenizer-ai-content')

text = st.text_area("Paste your Content (512 word limit)")

if text:
    tokenized_input = tokenizer(text, return_tensors='pt')
    with torch.no_grad():
        outputs = model(**tokenized_input)

    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=-1).item()
    st.write(f'Predicted Label: {labels_ids[predicted_class]}')