Renamed files and classes, updated README instructions to use pip

README.md

@@ -10,15 +10,8 @@ tags:
 
 # Electra Base Classifier for Sentiment Analysis
 
-This model is a fine-tuned [ELECTRA base model](https://huggingface.co/google/electra-base-discriminator) for sentiment analysis, enhanced with custom pooling and a SwishGLU activation function. It classifies text into three sentiment categories: negative, neutral, and positive.
+This is an [ELECTRA base discriminator](https://huggingface.co/google/electra-base-discriminator) fine-tuned for sentiment analysis of reviews. It has a mean pooling layer and a classifier head (2 layers of 1024 dimension) with SwishGLU activation and dropout (0.3). It classifies text into three sentiment categories: 'negative' (0), 'neutral' (1), and 'positive' (2). It was fine-tuned on the [Sentiment Merged](https://huggingface.co/datasets/jbeno/sentiment_merged) dataset, which is a merge of Stanford Sentiment Treebank (SST-3), and DynaSent Rounds 1 and 2.
 
-## Model Architecture
-
-- **Base Model**: ELECTRA base (`google/electra-base-discriminator`)
-- **Custom Components**:
-  - **Pooling Layer**: Custom pooling layer supporting 'cls', 'mean', and 'max' pooling types.
-  - **Activation Function**: Custom SwishGLU activation function.
-  - **Classifier**: Custom classifier with configurable hidden dimensions, number of layers, and dropout rate.
 
 ## Labels
 
@@ -30,34 +23,148 @@ The model predicts the following labels:
 
 ## How to Use
 
-To use this model, you need to download the `electra_base_classifier_sentiment.py` file from this repository and place it in your working directory. This file contains the custom classes required to load and use the model.
+### Install package
+
+This model requires the classes in `electra_classifier.py`. You can download the file, or you can install the package from PyPI.
+
+```bash
+pip install electra-classifier
+```
 
+### Load classes and model
 ```python
+# Install the package in a notebook
+!pip install electra-classifier
+
+# Import libraries
 import torch
 from transformers import AutoTokenizer
-from electra_base_classifier_sentiment import ElectraBaseClassifierSentiment  # Ensure this file is in your working directory
+from electra_classifier import ElectraClassifier
 
+# Load tokenizer and model
 model_name = "jbeno/electra-base-classifier-sentiment"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = ElectraBaseClassifierSentiment.from_pretrained(model_name)
+model = ElectraClassifier.from_pretrained(model_name)
+
+# Set model to evaluation mode
 model.eval()
 
-# Example usage
-text = "I love this product!"
+# Run inference
+text = "I love this restaurant!"
 inputs = tokenizer(text, return_tensors="pt")
+
 with torch.no_grad():
     logits = model(**inputs)
     predicted_class_id = torch.argmax(logits, dim=1).item()
-    predicted_label = model.config.id2label[str(predicted_class_id)]
+    predicted_label = model.config.id2label[predicted_class_id]
     print(f"Predicted label: {predicted_label}")
 ```
 
 ## Requirements
 - Python 3.7+
 - PyTorch
-- Transformers library
-- Additional Files:
-    - Download **electra_base_classifier_sentiment.py** from this repository and place it in your working directory.
+- Transformers
+- [electra-classifier](https://pypi.org/project/electra-classifier/) - Install with pip, or download electra_classifier.py
+
+## Training Details
+
+### Dataset
+
+The model was trained on the [Sentiment Merged](https://huggingface.co/datasets/jbeno/sentiment_merged) dataset, which is a mix of Stanford Sentiment Treebank (SST-3), DynaSent Round 1, and DynaSent Round 2.
+
+### Code
+
+The code used to train the model can be found on GitHub:
+- [jbeno/sentiment](https://github.com/jbeno/sentiment)
+- [jbeno/electra-classifier](https://github.com/jbeno/electra-classifier)
+
+### Research Paper
+
+The research paper can be found here: [ELECTRA and GPT-4o: Cost-Effective Partners for Sentiment Analysis](https://github.com/jbeno/sentiment/research_paper.pdf)
+
+### Performance
+
+- **Merged Dataset**
+    - Macro Average F1: **79.29**
+    - Accuracy: **79.69**
+- **DynaSent R1**
+    - Macro Average F1: **82.10**
+    - Accuracy: **82.14**
+- **DynaSent R2**
+    - Macro Average F1: **71.83**
+    - Accuracy: **71.94**
+- **SST-3**
+    - Macro Average F1: **69.95**
+    - Accuracy: **78.24**
+
+
+## Model Architecture
+
+- **Base Model**: ELECTRA base discriminator (`google/electra-base-discriminator`)
+- **Pooling Layer**: Custom pooling layer supporting 'cls', 'mean', and 'max' pooling types.
+- **Classifier**: Custom classifier with configurable hidden dimensions, number of layers, and dropout rate.
+    - **Activation Function**: Custom SwishGLU activation function.
+
+```
+ElectraClassifier(
+  (electra): ElectraModel(
+    (embeddings): ElectraEmbeddings(
+      (word_embeddings): Embedding(30522, 768, padding_idx=0)
+      (position_embeddings): Embedding(512, 768)
+      (token_type_embeddings): Embedding(2, 768)
+      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
+      (dropout): Dropout(p=0.1, inplace=False)
+    )
+    (encoder): ElectraEncoder(
+      (layer): ModuleList(
+        (0-11): 12 x ElectraLayer(
+          (attention): ElectraAttention(
+            (self): ElectraSelfAttention(
+              (query): Linear(in_features=768, out_features=768, bias=True)
+              (key): Linear(in_features=768, out_features=768, bias=True)
+              (value): Linear(in_features=768, out_features=768, bias=True)
+              (dropout): Dropout(p=0.1, inplace=False)
+            )
+            (output): ElectraSelfOutput(
+              (dense): Linear(in_features=768, out_features=768, bias=True)
+              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
+              (dropout): Dropout(p=0.1, inplace=False)
+            )
+          )
+          (intermediate): ElectraIntermediate(
+            (dense): Linear(in_features=768, out_features=3072, bias=True)
+            (intermediate_act_fn): GELUActivation()
+          )
+          (output): ElectraOutput(
+            (dense): Linear(in_features=3072, out_features=768, bias=True)
+            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
+            (dropout): Dropout(p=0.1, inplace=False)
+          )
+        )
+      )
+    )
+  )
+  (pooling): PoolingLayer()
+  (classifier): Classifier(
+    (layers): Sequential(
+      (0): Linear(in_features=768, out_features=1024, bias=True)
+      (1): SwishGLU(
+        (projection): Linear(in_features=1024, out_features=2048, bias=True)
+        (activation): SiLU()
+      )
+      (2): Dropout(p=0.3, inplace=False)
+      (3): Linear(in_features=1024, out_features=1024, bias=True)
+      (4): SwishGLU(
+        (projection): Linear(in_features=1024, out_features=2048, bias=True)
+        (activation): SiLU()
+      )
+      (5): Dropout(p=0.3, inplace=False)
+      (6): Linear(in_features=1024, out_features=3, bias=True)
+    )
+  )
+)
+```
+
 
 ## Custom Model Components
 
@@ -148,35 +255,6 @@ The model's configuration (config.json) includes custom parameters:
 - `dropout_rate`: Dropout rate used in the classifier.
 - `pooling`: Pooling strategy used ('mean').
 
-## Training Details
-
-### Dataset
-
-The model was trained on the [Sentiment Merged](https://huggingface.co/datasets/jbeno/sentiment_merged) dataset, which is a mix of Stanford Sentiment Treebank (SST-3), DynaSent Round 1, and DynaSent Round 2.
-
-### Code
-
-The code used to train the model can be found on GitHub: [jbeno/sentiment](https://github.com/jbeno/sentiment)
-
-### Research Paper
-
-The research paper can be found here: [ELECTRA and GPT-4o: Cost-Effective Partners for Sentiment Analysis](https://github.com/jbeno/sentiment/research_paper.pdf)
-
-### Performance
-
-- **Merged Dataset**
-    - Macro Average F1: **79.29**
-    - Accuracy: **79.69**
-- **DynaSent R1**
-    - Macro Average F1: **82.10**
-    - Accuracy: **82.14**
-- **DynaSent R2**
-    - Macro Average F1: **71.83**
-    - Accuracy: **71.94**
-- **SST-3**
-    - Macro Average F1: **69.95**
-    - Accuracy: **78.24**
-
 ## License
 
 This model is licensed under the MIT License.

__init__.py

@@ -1 +1 @@
-from .electra_base_classifier_sentiment import ElectraBaseClassifierSentiment
+from .electra_classifier import ElectraClassifier

config.json

@@ -1,6 +1,6 @@
 {
   "architectures": [
-    "ElectraBaseClassifierSentiment"
+    "ElectraClassifier"
   ],
   "attention_probs_dropout_prob": 0.1,
   "classifier_dropout": null,

electra_base_classifier_sentiment.py → electra_classifier.py

@@ -1,23 +1,21 @@
 import torch
-import torch.nn as nn
+from torch import nn
 from transformers import ElectraPreTrainedModel, ElectraModel
 
+# Custom activation function
 class SwishGLU(nn.Module):
     def __init__(self, input_dim: int, output_dim: int):
         super(SwishGLU, self).__init__()
-        # Linear projection to 2 * output_dim to split for gate and projection
         self.projection = nn.Linear(input_dim, 2 * output_dim)
-        self.activation = nn.SiLU()  # Swish activation
+        self.activation = nn.SiLU()
 
     def forward(self, x):
-        # Apply linear projection
         x_proj_gate = self.projection(x)
-        # Split the projection into two parts
         projected, gate = x_proj_gate.tensor_split(2, dim=-1)
-        # Apply Swish activation and multiply
         return projected * self.activation(gate)
 
-    
+
+# Custom pooling layer
 class PoolingLayer(nn.Module):
     def __init__(self, pooling_type='cls'):
         super().__init__()
@@ -35,6 +33,8 @@ class PoolingLayer(nn.Module):
         else:
             raise ValueError(f"Unknown pooling method: {self.pooling_type}")
 
+
+# Custom classifier
 class Classifier(nn.Module):
     def __init__(self, input_dim, hidden_dim, hidden_activation, num_layers, n_classes, dropout_rate=0.0):
         super().__init__()
@@ -56,7 +56,9 @@ class Classifier(nn.Module):
     def forward(self, x):
         return self.layers(x)
 
-class ElectraBaseClassifierSentiment(ElectraPreTrainedModel):
+
+# Custom Electra classifier model
+class ElectraClassifier(ElectraPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.electra = ElectraModel(config)
@@ -87,7 +89,6 @@ class ElectraBaseClassifierSentiment(ElectraPreTrainedModel):
         )
         self.init_weights()
 
-
     def forward(self, input_ids=None, attention_mask=None, **kwargs):
         outputs = self.electra(input_ids, attention_mask=attention_mask, **kwargs)
         pooled_output = self.pooling(outputs.last_hidden_state, attention_mask)