Initial commit

accessibility_classifier/api.py

@@ -0,0 +1,11 @@
+from fastapi import Depends, FastAPI
+from .classifier.model import Model, get_model
+
+app = FastAPI()
+
+
+@app.post("/predict")
+async def predict(issue: str, model: Model = Depends(get_model)):
+    accessibility_confidence, nonaccessibility_confidence = model.predict(issue)
+    return {"Accessibility": str(accessibility_confidence),
+            "Non-accessibility": str(nonaccessibility_confidence)}

accessibility_classifier/classifier/model.py

@@ -0,0 +1,115 @@
+from transformers.modeling_outputs import SequenceClassifierOutput
+from transformers import AlbertForSequenceClassification, AlbertTokenizer
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+class AlbertForMultilabelSequenceClassification(AlbertForSequenceClassification):
+    def __init__(self, config):
+        super().__init__(config)
+
+    def forward(self,
+                input_ids=None,
+                attention_mask=None,
+                token_type_ids=None,
+                position_ids=None,
+                head_mask=None,
+                inputs_embeds=None,
+                labels=None,
+                output_attentions=None,
+                output_hidden_states=None,
+                return_dict=None):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.albert(input_ids,
+                              attention_mask=attention_mask,
+                              token_type_ids=token_type_ids,
+                              position_ids=position_ids,
+                              head_mask=head_mask,
+                              inputs_embeds=inputs_embeds,
+                              output_attentions=output_attentions,
+                              output_hidden_states=output_hidden_states,
+                              return_dict=return_dict)
+
+        pooled_output = outputs[1]
+        pooled_output = self.dropout(pooled_output)
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = torch.nn.BCEWithLogitsLoss()
+            loss = loss_fct(logits.view(-1, self.num_labels),
+                            labels.float().view(-1, self.num_labels))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(loss=loss,
+                                        logits=logits,
+                                        hidden_states=outputs.hidden_states,
+                                        attentions=outputs.attentions)
+
+
+class Model:
+    def __init__(self):
+
+        self.device = torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu")
+
+        self.labels = ['Accessibility', 'Non-accessibility']
+        self.tokenizer = AlbertTokenizer.from_pretrained(
+            'albert-base-v2', do_lower_case=True)
+        classifier = AlbertForMultilabelSequenceClassification.from_pretrained(
+            'albert-base-v2',
+            output_attentions=False,
+            output_hidden_states=False,
+            num_labels=2
+        )
+        
+        classifier.load_state_dict(
+            torch.load("assets/pytorch_model.bin", map_location=self.device))
+        classifier = classifier.eval()
+        self.classifier = classifier.to(self.device)
+
+    def predict(self, text):
+        encoded_text = self.tokenizer.encode_plus(
+            text,
+            max_length=30,
+            add_special_tokens=True,
+            return_token_type_ids=False,
+            padding='longest',
+            return_attention_mask=True,
+            return_tensors="pt",
+            truncation=True,
+
+        )
+
+        input_ids = encoded_text["input_ids"].to(self.device)
+        attention_mask = encoded_text["attention_mask"].to(self.device)
+
+        with torch.no_grad():
+            probabilities = self.classifier(input_ids, attention_mask)
+
+        prediction = F.softmax(probabilities.logits,
+                               dim=1).cpu().numpy().flatten().max()
+        prediction_index = np.where(F.softmax(probabilities.logits,
+                                              dim=1).cpu().numpy() == prediction)[1][0]
+        label = self.labels[prediction_index]
+
+        all_predictions = F.softmax(
+            probabilities.logits, dim=1).cpu().numpy().flatten()
+
+        accessibility_prediction = all_predictions[0]
+        nonaccessibility_prediction = all_predictions[1]
+
+        return (accessibility_prediction, nonaccessibility_prediction)
+
+
+model = Model()
+# model.predict("this is an impsorvement")
+
+
+def get_model():
+    return model

assets/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75b432cd7171f3d2ff0150c24641ce9b956165f5a73f439a1c0b56d4f963d1bf
+size 46749762

bin/download_model

@@ -0,0 +1,6 @@
+# !/usr/bin/env python
+import gdown
+
+gdown.download(
+    "assets/pytorch_model.bin",
+)

bin/start_server

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+uvicorn --port 8081 accessibility_classifier.api:app

bin/test_request

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+http POST http://localhost:8000/predict text="This is an accessibility issue"