Update pages/imdb.py

pages/imdb.py

@@ -8,6 +8,8 @@ import re
 import string
 from collections import Counter
 
+from transformers import AutoConfig, AutoTokenizer, AutoModel, AutoModelForSequenceClassification, Trainer, TrainingArguments
+
 from gensim.models import Word2Vec
 from string import punctuation
 import transformers
@@ -138,12 +140,10 @@ def main():
 
     start1 = time.time()
     
-    automodel = transformers.AutoModelForSequenceClassification.from_pretrained(
-    'distilbert-base-uncased-finetuned-sst-2-english'
-    )
-    autotoken = transformers.AutoTokenizer.from_pretrained(
-    'distilbert-base-uncased-finetuned-sst-2-english'
-    )
+    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+    config = AutoConfig.from_pretrained('distilbert-base-uncased', num_labels=2)
+    
+    automodel = AutoModelForSequenceClassification.from_config(config)
     
 
     input_tokens = autotoken(