Update model.py

model.py

@@ -10,7 +10,7 @@ Original file is located at
 """
 
 import torch
-from transformers import AutoTokenizer, BertForSequenceClassification
+from transformers import AutoTokenizer, BertForSequenceClassification,AutoConfig
 from datasets import Dataset
 import pandas as pd
 import os
@@ -25,8 +25,9 @@ stopwords = stopwords_list()
 # Load the BERT model for sentiment analysis
 dataset = Dataset.from_pandas(pd.DataFrame({"Comment": []}))
 
-tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-base-parsbert-uncased")
-model = BertForSequenceClassification.from_pretrained("HooshvareLab/bert-base-parsbert-uncased", num_labels=3)
+config = AutoConfig.from_pretrained("HooshvareLab/albert-fa-zwnj-base-v2")
+tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/albert-fa-zwnj-base-v2")
+model = BertForSequenceClassification.from_pretrained("HooshvareLab/albert-fa-zwnj-base-v2", num_labels=3)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
@@ -48,7 +49,7 @@ def predict_sentiment(batch):
     return {'sentiment': predictions.cpu()}
 
 # Mapping sentiment labels
-sentiment_labels_en = {0: 'منفی', 1: 'خنثی', 2: 'مثبت'}
+#sentiment_labels_en = {0: 'منفی', 1: 'خنثی', 2: 'مثبت'}
 
 
 # Adding sentiment prediction to tokenized dataset