update app.py

app.py

@@ -1,10 +1,13 @@
 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")
+text = st.text_area("Paste your Content (512 word limit)")
 
 if text:
     tokenized_input = tokenizer(text, return_tensors='pt')
@@ -13,4 +16,4 @@ if text:
 
     logits = outputs.logits
     predicted_class = torch.argmax(logits, dim=-1).item()
-    st.json(predicted_class)
+    st.write(f'Predicted Label: {labels_ids[predicted_class]}')