Deploying Roberta Sentinel

ModelDriver.py

@@ -0,0 +1,54 @@
+from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaModel
+import torch
+import torch.nn as nn
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+class MLP(nn.Module):
+    def __init__(self, input_dim):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_dim, 256)
+        self.fc2 = nn.Linear(256, 2)
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        x = self.gelu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+def extract_features(text):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+    model = RobertaModel.from_pretrained("roberta-base").to(device)
+    tokenized_text = tokenizer.encode(text, truncation=True, max_length=512, return_tensors="pt")
+    outputs = model(tokenized_text)
+    last_hidden_states = outputs.last_hidden_state
+    TClassification = last_hidden_states[:, 0, :].squeeze().detach().numpy()
+    return TClassification
+
+def RobertaSentinelOpenGPTInference(input_text):
+    features = extract_features(input_text)
+    loaded_model = MLP(768).to(device)
+    loaded_model.load_state_dict(torch.load("MLPDictStates/RobertaSentinelOpenGPT.pth"))
+
+    # Define the tokenizer and model for feature extraction
+    with torch.no_grad():
+        inputs = torch.tensor(features).to(device)
+        outputs = loaded_model(inputs.float())
+        _, predicted = torch.max(outputs, 1)
+
+    return predicted.item()
+
+def RobertaSentinelCSAbstractInference(input_text):
+    features = extract_features(input_text)
+    loaded_model = MLP(768).to(device)
+    loaded_model.load_state_dict(torch.load("MLPDictStates/RobertaSentinelCSAbstract.pth"))
+
+    # Define the tokenizer and model for feature extraction
+    with torch.no_grad():
+        inputs = torch.tensor(features).to(device)
+        outputs = loaded_model(inputs.float())
+        _, predicted = torch.max(outputs, 1)
+
+    return predicted.item()
+
+

app.py

@@ -1,25 +1,26 @@
 import streamlit as st
 from transformers import pipeline
+from ModelDriver import RobertaSentinelOpenGPTInference, RobertaSentinelCSAbstractInference
 
 # Add a title
 st.title('GPT Detection Demo')
 
 # Add 4 options for 4 models
 option = st.sidebar.selectbox(
-    'Which pipeline do you want to use?',
-    ('sentiment-analysis', 'ner', 'question-answering', 'text-generation'),
+    'Which Model do you want to use?',
+    ('RobertaSentinelOpenGPT', 'RobertaSentinelCSAbstract'),
 )
 
-option2 = st.sidebar.selectbox(
-    'Which model do you want to use?',
-    ('gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'),
-)
 
-pipe = pipeline(option)
 text = st.text_area('Enter text here', '')
 
 if st.button('Generate'):
-    st.write(pipe(text)[0])
+    if option == 'RobertaSentinelOpenGPT':
+        result = RobertaSentinelOpenGPTInference(text)
+    elif option == 'RobertaSentinelCSAbstract':
+        result = RobertaSentinelCSAbstractInference(text)
+    st.write(result)
+