test11

Dockerfile

@@ -9,8 +9,6 @@ COPY ./requirements.txt /code/requirements.txt
 
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
-COPY ~/.cache/torch/sentence_transformers /root/.cache/torch/sentence_transformers
-
 COPY . .
 
 CMD ["gunicorn", "-b", "0.0.0.0:7860", "main:app"]

main.py

@@ -1,29 +1,7 @@
 from flask import Flask, request
 import torch
-import torch.nn as nn
+import model
 import numpy as np
-from sentence_transformers import SentenceTransformer
-
-class LSTM(nn.Module):
-    def __init__(self, embedding_dim, hidden_dim, num_layers, output_dim):
-        super(LSTM, self).__init__()
-        self.lstm1 = nn.LSTM(embedding_dim, hidden_dim, num_layers, batch_first=True)
-        self.lstm2 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
-        self.lstm3 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
-        self.lstm4 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
-        self.lstm5 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
-        self.o = nn.Linear(hidden_dim, output_dim)
-
-    def forward(self, embedding):
-        o_n1, (h_n1, c_n1) = self.lstm1(embedding)
-        o_n2, (h_n2, c_n2) = self.lstm2(o_n1, (h_n1, c_n1))
-        o_n3, (h_n3, c_n3) = self.lstm3(o_n2, (h_n2, c_n2))
-        o_n4, (h_n4, c_n4) = self.lstm4(o_n3, (h_n3, c_n3))
-        o_n5, (h_n5, c_n5) = self.lstm5(o_n4, (h_n4, c_n4))
-        output = self.o(o_n5)
-        return output
-
-sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
 
 embedding_dim = 384
 hidden_dim = 512
@@ -32,14 +10,14 @@ output_dim = 180
 num_epochs = 100
 learning_rate = 0.001
 
-lstm_model = LSTM(embedding_dim, hidden_dim, num_layers, output_dim)
+lstm_model = model.LSTM(embedding_dim, hidden_dim, num_layers, output_dim)
 lstm_model.load_state_dict(torch.load('lstm.pt'))
 
 app = Flask(__name__)
 
 def GeneratePosesJSON(input):
     with torch.no_grad():
-        processed_text = torch.tensor(sentence_model.encode(input), dtype=torch.float)
+        processed_text = torch.tensor(input, dtype=torch.float)
         output_poses = lstm_model(processed_text.unsqueeze(0))
 
         people = output_poses.cpu().detach().numpy().reshape(5, 18, 2).tolist()

model.py

@@ -0,0 +1,20 @@
+import torch.nn as nn
+
+class LSTM(nn.Module):
+    def __init__(self, embedding_dim, hidden_dim, num_layers, output_dim):
+        super(LSTM, self).__init__()
+        self.lstm1 = nn.LSTM(embedding_dim, hidden_dim, num_layers, batch_first=True)
+        self.lstm2 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
+        self.lstm3 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
+        self.lstm4 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
+        self.lstm5 = nn.LSTM(hidden_dim, hidden_dim, num_layers, batch_first=True)
+        self.o = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, embedding):
+        o_n1, (h_n1, c_n1) = self.lstm1(embedding)
+        o_n2, (h_n2, c_n2) = self.lstm2(o_n1, (h_n1, c_n1))
+        o_n3, (h_n3, c_n3) = self.lstm3(o_n2, (h_n2, c_n2))
+        o_n4, (h_n4, c_n4) = self.lstm4(o_n3, (h_n3, c_n3))
+        o_n5, (h_n5, c_n5) = self.lstm5(o_n4, (h_n4, c_n4))
+        output = self.o(o_n5)
+        return output