Update app.py

app.py

@@ -3,13 +3,18 @@ import torch.nn as nn
 import matplotlib.pyplot as plt
 import streamlit as st
 
-# Simple LSTM with a VRAM-like 4x4 Memory Grid
+# Simple LSTM with a VRAM-like 4x4 Memory Grid, including gates and filters
 class MemristorLSTM(nn.Module):
     def __init__(self, memory_size=4):
         super(MemristorLSTM, self).__init__()
         self.memory_size = memory_size
         self.lstm = nn.LSTM(input_size=1, hidden_size=50, num_layers=2, batch_first=True)
-        self.fc = nn.Linear(50, 1)
+        self.fc = nn.Linear(50, self.memory_size * self.memory_size)  # Output 16 values
+        # Gate to modulate the output
+        self.gate = nn.Sigmoid()
+        # Tanh filter to apply non-linearity to the output
+        self.filter = nn.Tanh()
+
         # 4x4 memory grid initialized to zeros
         self.memory = torch.zeros(memory_size, memory_size)  
 
@@ -18,9 +23,15 @@ class MemristorLSTM(nn.Module):
         lstm_out, _ = self.lstm(x)
         output = self.fc(lstm_out[:, -1, :])
 
+        # Apply sigmoid gate to control the flow of information
+        gated_output = self.gate(output)
+
+        # Apply tanh filter for non-linearity
+        filtered_output = self.filter(gated_output)
+
         # Update memory (4x4 grid), simulating synaptic memory flow
-        self.memory = self.memory + output.view(self.memory_size, self.memory_size)
-        return output, self.memory
+        self.memory = self.memory + filtered_output.view(self.memory_size, self.memory_size)
+        return filtered_output, self.memory
 
 # Load Pretrained Model
 def load_model(model_path="memristor_lstm.pth"):
@@ -62,7 +73,7 @@ def generate_spikes_for_pressure(pressure):
         visualize_memory_grid(memory_grid)
 
         # Optionally, print the output
-        print(f"Output: {output}")
+        print(f"Filtered Output: {output}")
         print(f"Memory Grid after update:\n{memory_grid}")
 
         return output, memory_grid