Update README.md

README.md

@@ -46,22 +46,28 @@ The main highlight of this model is the integration of a rare **bilinear layer (
 You can download the architecture file and the model weights directly from this repository:
 
 ```python
-import torch as t
-from huggingface_hub import hf_hub_download
+class WebAISC(nn.Module):
 
-# 1. Download the architecture and weights files
-REPO_ID = "ALEXFLR/sbl-net"
-hf_hub_download(repo_id=REPO_ID, filename="model.py", local_dir=".")
-hf_hub_download(repo_id=REPO_ID, filename="model_weights_hybrid.pth", local_dir=".")
+    def __init__(self):
+        super().__init__()
+        self.bilinear = nn.Bilinear(in1_features=2, in2_features=2, out_features=250)
 
-# 2. Import the model class and load the weights
-from model import WebAISC
+        self.x2 = nn.Linear(250, 100)
+        self.x3 = nn.Linear(100, 250)
+        self.x4 = nn.Linear(250, 1)
+        self.selu = nn.SELU()
 
-model = WebAISC()
-model.load_state_dict(t.load("model_weights_hybrid.pth", map_location=t.device('cpu')))
-model.eval()
+    def forward(self, x):
+        a = x[:, 0:2]
+        b = x[:, 2:4]
+
+        x = self.selu(self.bilinear(a, b))
+
+        x = self.selu(self.x2(x))
+        x = self.selu(self.x3(x))
+        x = self.x4(x)
+        return x
 
-# 3. Inference
 test_input = t.tensor([[-6.0, 70.0, 4.0, -196.0]], dtype=t.float32)
 with t.no_grad():
     prediction = model(test_input)