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

qlib/contrib/model/pytorch_sfm.py

@@ -7,10 +7,9 @@ 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,
-    save_multiple_parts_file,
     get_or_create_path,
     drop_nan_by_y_index,
 )
@@ -442,11 +441,11 @@ class SFM(Model):
 
         raise ValueError("unknown metric `%s`" % self.metric)
 
-    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.sfm_model.eval()
         x_values = x_test.values
@@ -459,10 +458,7 @@ class SFM(Model):
             else:
                 end = begin + self.batch_size
 
-            x_batch = torch.from_numpy(x_values[begin:end]).float()
-
-            if self.device != "cpu":
-                x_batch = x_batch.to(self.device)
+            x_batch = torch.from_numpy(x_values[begin:end]).float().to(self.device)
 
             with torch.no_grad():
                 pred = self.sfm_model(x_batch).detach().cpu().numpy()