Update app.py

app.py

@@ -1,17 +1,25 @@
 from flask import Flask, request
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import RobertaConfig
 from transformers import RobertaForSequenceClassification, RobertaTokenizer, RobertaConfig
 import torch
+from torch import cuda
 import gradio as gr
 import os
 import re
 app = Flask(__name__)
 
 ACCESS_TOKEN = os.environ["ACCESS_TOKEN"]
-config = RobertaConfig.from_pretrained("PirateXX/ChatGPT-Text-Detector", use_auth_token= ACCESS_TOKEN)
-model = RobertaForSequenceClassification.from_pretrained("PirateXX/ChatGPT-Text-Detector", use_auth_token= ACCESS_TOKEN, config = config)
+# config = RobertaConfig.from_pretrained("PirateXX/ChatGPT-Text-Detector", use_auth_token= ACCESS_TOKEN)
+# model = RobertaForSequenceClassification.from_pretrained("PirateXX/ChatGPT-Text-Detector", use_auth_token= ACCESS_TOKEN, config = config)
 
-model_name = "roberta-base"
-tokenizer = RobertaTokenizer.from_pretrained(model_name, map_location=torch.device('cpu'))
+device = 'cuda' if cuda.is_available() else 'cpu'
+tokenizer = AutoTokenizer.from_pretrained("PirateXX/AI-Content-Detector", use_auth_token= "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")
+model = AutoModelForSequenceClassification.from_pretrained("PirateXX/AI-Content-Detector", use_auth_token= "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")
+model.to(device)
+
+# model_name = "roberta-base"
+# tokenizer = RobertaTokenizer.from_pretrained(model_name, map_location=torch.device('cpu'))
 
 def text_to_sentences(text):
     clean_text = text.replace('\n', ' ')
@@ -34,7 +42,7 @@ def chunks_of_900(text, chunk_size=900):
     chunks.append(current_chunk)
     return chunks
     
-def predict(query, device="cpu"):
+def predict(query):
     tokens = tokenizer.encode(query)
     all_tokens = len(tokens)
     tokens = tokens[:tokenizer.model_max_length - 2]