Update model.py

model.py

@@ -1,20 +1,15 @@
 import torch
-import tensorflow as tf
-from transformers import BertTokenizer, BertForSequenceClassification, TFAutoModel
+from transformers import BertTokenizer, BertForSequenceClassification
 
 
 def get_model():
-    model = TFAutoModel.from_pretrained('bert-base-uncased')
+    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
     return model
 
 
 # Predicting Function
 def predict(model, text, tokenizer):
-    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="tf")
+    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
     outputs = model(**inputs)
-    predictions = tf.argmax(outputs.logits, dim=-1)
-    return "AI-generated" if predictions.numpy().item() == 1 else "Human-written"
-
-# Example Usage (commented out as it's not needed for web deployment)
-# user_input = input("Enter the text you want to classify: ")
-# print("Classified as:", predict(user_input))
+    predictions = torch.argmax(outputs.logits, dim=-1)
+    return "AI-generated" if predictions.item() == 1 else "Human-written"