Add TREC classifier code and requirements

app.py

@@ -0,0 +1,227 @@
+import torch
+import torch.nn as nn
+import lightning as L
+import torchmetrics as tm
+from tokenizers import Tokenizer
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+COARSE_LABELS = [
+    "ABBR (0): Abbreviation",
+    "ENTY (1): Entity",
+    "DESC (2): Description and abstract concept",
+    "HUM (3): Human being",
+    "LOC (4): Location",
+    "NUM (5): Numeric value",
+]
+
+FINE_LABELS = [
+    "ABBR (0): Abbreviation",
+    "ABBR (1): Expression abbreviated",
+    "ENTY (2): Animal",
+    "ENTY (3): Organ of body",
+    "ENTY (4): Color",
+    "ENTY (5): Invention, book and other creative piece",
+    "ENTY (6): Currency name",
+    "ENTY (7): Disease and medicine",
+    "ENTY (8): Event",
+    "ENTY (9): Food",
+    "ENTY (10): Musical instrument",
+    "ENTY (11): Language",
+    "ENTY (12): Letter like a-z",
+    "ENTY (13): Other entity",
+    "ENTY (14): Plant",
+    "ENTY (15): Product",
+    "ENTY (16): Religion",
+    "ENTY (17): Sport",
+    "ENTY (18): Element and substance",
+    "ENTY (19): Symbols and sign",
+    "ENTY (20): Techniques and method",
+    "ENTY (21): Equivalent term",
+    "ENTY (22): Vehicle",
+    "ENTY (23): Word with a special property",
+    "DESC (24): Definition of something",
+    "DESC (25): Description of something",
+    "DESC (26): Manner of an action",
+    "DESC (27): Reason",
+    "HUM (28): Group or organization of persons",
+    "HUM (29): Individual",
+    "HUM (30): Title of a person",
+    "HUM (31): Description of a person",
+    "LOC (32): City",
+    "LOC (33): Country",
+    "LOC (34): Mountain",
+    "LOC (35): Other location",
+    "LOC (36): State",
+    "NUM (37): Postcode or other code",
+    "NUM (38): Number of something",
+    "NUM (39): Date",
+    "NUM (40): Distance, linear measure",
+    "NUM (41): Price",
+    "NUM (42): Order, rank",
+    "NUM (43): Other number",
+    "NUM (44): Lasting time of something",
+    "NUM (45): Percent, fraction",
+    "NUM (46): Speed",
+    "NUM (47): Temperature",
+    "NUM (48): Size, area and volume",
+    "NUM (49): Weight",
+]
+
+
+class Classifier:
+    def __init__(self, tokenizer_ckpt_path, model_ckpt_path):
+        self.tokenizer = Tokenizer.from_file(tokenizer_ckpt_path)
+        self.model = LSTMWithAttentionClassifier.load_from_checkpoint(
+            model_ckpt_path,
+            map_location="cpu",
+        )
+
+    def predict(self, text):
+        encoding = self.tokenizer.encode(text)
+        ids = torch.tensor([encoding.ids])
+        logits, _ = self.model(ids)
+        probs = torch.softmax(logits, dim=1).squeeze().tolist()
+        return {
+            category: prob
+            for category, prob in zip(
+                FINE_LABELS if self.model.fine else COARSE_LABELS, probs
+            )
+        }
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden_dim):
+        super().__init__()
+        self.WQuery = nn.Linear(hidden_dim, hidden_dim)
+        self.WKey = nn.Linear(hidden_dim, hidden_dim)
+        self.WValue = nn.Linear(hidden_dim, 1)
+
+    def forward(self, x):
+        query = torch.tanh(self.WQuery(x))
+        key = torch.tanh(self.WKey(x))
+
+        attention_weights = torch.softmax(self.WValue(query + key), dim=1)
+
+        return (attention_weights * x).sum(dim=1), attention_weights
+
+
+class LSTMWithAttentionClassifier(L.LightningModule):
+    def __init__(
+        self,
+        vocab_size,
+        embedding_dim,
+        hidden_dim,
+        num_classes,
+        lr=1e-3,
+        weight_decay=1e-2,
+        num_layers=1,
+        bidirectional=False,
+        dropout=0.0,
+        padding_idx=3,
+        fine=False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.fine = fine
+
+        self.embedding = nn.Embedding(
+            vocab_size,
+            embedding_dim,
+            padding_idx=padding_idx,
+        )
+        self.lstm = nn.LSTM(
+            embedding_dim,
+            hidden_dim,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+            dropout=dropout,
+        )
+        self.attention = Attention(
+            hidden_dim * (1 + bidirectional),
+        )
+        self.fc = nn.Linear(
+            hidden_dim * (1 + bidirectional),
+            num_classes,
+        )
+
+        self.criteria = nn.CrossEntropyLoss()
+        self.accuracy = tm.Accuracy(
+            task="multiclass",
+            num_classes=num_classes,
+        )
+
+    def forward(self, input_ids):
+        x = self.embedding(input_ids)
+        x, _ = self.lstm(x)
+        x, attention_weights = self.attention(x)
+        x = self.fc(x)
+        return x, attention_weights
+
+    def training_step(self, batch, batch_idx):
+        input_ids = batch["input_ids"]
+        coarse = batch["coarse"]
+        fine = batch["fine"]
+        logits, _ = self(input_ids)
+        loss = self.criteria(logits, fine if self.fine else coarse)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        input_ids = batch["input_ids"]
+        coarse = batch["coarse"]
+        fine = batch["fine"]
+        logits, _ = self(input_ids)
+        loss = self.criteria(logits, fine if self.fine else coarse)
+        self.log("val_loss", loss)
+        pred = logits.argmax(dim=1)
+        self.accuracy(pred, fine if self.fine else coarse)
+        self.log("val_acc", self.accuracy, prog_bar=True)
+
+    def configure_optimizers(self):
+        return torch.optim.AdamW(
+            self.parameters(),
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+        )
+
+
+tokenizer_ckpt_path = hf_hub_download(
+    repo_id="SatwikKambham/trec-classifier",
+    filename="tokenizer.json",
+)
+model_ckpt_path = hf_hub_download(
+    repo_id="SatwikKambham/trec-classifier",
+    filename="lstm_attention.ckpt",
+)
+classifier = Classifier(tokenizer_ckpt_path, model_ckpt_path)
+interface = gr.Interface(
+    fn=classifier.predict,
+    inputs=gr.components.Textbox(
+        label="Question",
+        placeholder="Enter a question here...",
+    ),
+    outputs=gr.components.Label(
+        label="Predicted class",
+        num_top_classes=3,
+    ),
+    examples=[
+        [
+            "What does LOL mean?",
+        ],
+        [
+            "What is the meaning of life?",
+        ],
+        [
+            "How long does it take for light from the sun to reach the earth?",
+        ],
+        [
+            "When is friendship day?",
+        ],
+    ],
+)
+interface.launch()

requirements.txt

@@ -0,0 +1,5 @@
+torch
+tokenizers
+lightning
+torchmetrics
+huggingface_hub