import json

from rest_framework import status, viewsets
from rest_framework.parsers import JSONParser
from rest_framework.response import Response

from .utils import predict


class PredictViewSet(viewsets.ViewSet):
    parser_classes = (JSONParser,)

    def create(self, request):
        """
        example call:
        curl -X POST "http://localhost:8000/" -H "content-type: application/json" --data '{"source": "1/5/2003"}'
        """
        if isinstance(request.data, str):
            data = json.loads(request.data)
        else:
            data = request.data

        if not "source" in data:
            return Response(
                'Format: {"source": string}', status=status.HTTP_400_BAD_REQUEST
            )
        return Response(predict(data.get("source")))

from allennlp.common.util import import_submodules
from allennlp.models.archival import load_archive
from allennlp.predictors import Predictor


def predict(source):
    # register custom model code
    # the library code is in the same Django app (main) in the folder "library"
    import_submodules("main.library")

    # load model without using cuda
    archive = load_archive("../models/model.tar.gz", cuda_device=-1)
    # select predictor needed for this model
    predictor = Predictor.from_archive(archive, "seq2seq")

    result = predictor.predict_json({"source": source})
    # add a confidence bool to help ignoring bad results
    confident = True if abs(result["class_log_probabilities"][0]) < 0.05 else False

    return {
        "predicted": "".join(result["predicted_tokens"]),
        "probability": result["class_log_probabilities"][0],
        "confident": confident,
    }

diff --git a/simple_seq2seq.py b/simple_seq2seq.py
index 849da8a..9c7e3da 100644
--- a/simple_seq2seq.py
+++ b/simple_seq2seq.py
@@ -323,6 +323,7 @@ class SimpleSeq2Seq(Model):

         step_logits: List[torch.Tensor] = []
         step_predictions: List[torch.Tensor] = []
+        attn: List[torch.Tensor] = []
         for timestep in range(num_decoding_steps):
             if self.training and torch.rand(1).item() < self._scheduled_sampling_ratio:
                 # Use gold tokens at test time and at a rate of 1 - _scheduled_sampling_ratio
@@ -338,6 +339,8 @@ class SimpleSeq2Seq(Model):

             # shape: (batch_size, num_classes)
             output_projections, state = self._prepare_output_projections(input_choices, state)
+            if not self.training:
+                attn.append(torch.squeeze(state["attention_weights"]))

             # list of tensors, shape: (batch_size, 1, num_classes)
             step_logits.append(output_projections.unsqueeze(1))
@@ -358,6 +361,9 @@ class SimpleSeq2Seq(Model):

         output_dict = {"predictions": predictions}

+        if not self.training:
+            output_dict["attentions"] = torch.unsqueeze(torch.stack(attn), 0)
+
         if target_tokens:
             # shape: (batch_size, num_decoding_steps, num_classes)
             logits = torch.cat(step_logits, 1)
@@ -412,7 +418,8 @@ class SimpleSeq2Seq(Model):

         if self._attention:
             # shape: (group_size, encoder_output_dim)
-            attended_input = self._prepare_attended_input(decoder_hidden, encoder_outputs, source_mask)
+            attended_input, input_weights = self._prepare_attended_input(decoder_hidden, encoder_outputs, source_mask)
+            state["attention_weights"] = input_weights

             # shape: (group_size, decoder_output_dim + target_embedding_dim)
             decoder_input = torch.cat((attended_input, embedded_input), -1)
@@ -451,7 +458,7 @@ class SimpleSeq2Seq(Model):
         # shape: (batch_size, encoder_output_dim)
         attended_input = util.weighted_sum(encoder_outputs, input_weights)

-        return attended_input
+        return attended_input, input_weights

     @staticmethod
     def _get_loss(logits: torch.LongTensor,

allennlp train configurations/reverse_starting_point.json -s output --include-package library

allennlp predict output/model.tar.gz --use-dataset-reader examples.csv --predictor my_seq2seq --output-file output/examples.output --include-package library
python tools/plot_attention.py