Collect all contrib models in __init__ and add unit tests for init

qlib/contrib/model/pytorch_lstm.py

@@ -8,6 +8,7 @@ from __future__ import print_function
 import os
 import numpy as np
 import pandas as pd
+from typing import Text, Union
 import copy
 from ...utils import (
     unpack_archive_with_buffer,
@@ -268,11 +269,11 @@ class LSTM(Model):
         if self.use_gpu:
             torch.cuda.empty_cache()
 
-    def predict(self, dataset):
+    def predict(self, dataset: DatasetH, segment: Union[Text, slice] = "test"):
         if not self.fitted:
             raise ValueError("model is not fitted yet!")
 
-        x_test = dataset.prepare("test", col_set="feature")
+        x_test = dataset.prepare(segment, col_set="feature", data_key=DataHandlerLP.DK_I)
         index = x_test.index
         self.lstm_model.eval()
         x_values = x_test.values
@@ -280,17 +281,13 @@ class LSTM(Model):
         preds = []
 
         for begin in range(sample_num)[:: self.batch_size]:
-
             if sample_num - begin < self.batch_size:
                 end = sample_num
             else:
                 end = begin + self.batch_size
-
             x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device)
-
             with torch.no_grad():
                 pred = self.lstm_model(x_batch).detach().cpu().numpy()
-
             preds.append(pred)
 
         return pd.Series(np.concatenate(preds), index=index)