Update app.py

app.py

@@ -61,7 +61,7 @@ def load_model(model_name: str):
     
         return model, tokenizer
 
-    elif model_name == 'CNN':
+    elif model_name == 'CNN-m2-8k-context':
         hidden_dim = 2048
         width = 2048
         seq_drop_prob = 0.05
@@ -84,6 +84,36 @@ def load_model(model_name: str):
         model.load_state_dict(weights)
         return model, None
 
+    elif model_name == 'CNN-m4-16k-context':
+        seq_drop_prob = 0.05
+        hidden_dim = 2000
+        width = 768
+        train_sequence_length = 16000
+        weight_decay = 0.0001
+        num_labs = len(set(y_train))
+        
+        
+        model_seq = nn.Sequential(
+            nn.Conv1d(4, width, 7, padding=3),
+            nn.ReLU(),
+            nn.BatchNorm1d(width),
+            ResNet1d(width, [(3, width // 2, 1)] * 1, dropout=None, dilation=7),
+            nn.ReLU(),
+            Pool2BN(width),
+        )
+        
+        new_head = torch.nn.Sequential(
+            torch.nn.Dropout(0.5),  ## for DEEPLIFT comment out
+            MLP([width * 2, num_labs])
+        )
+        
+        joined_model = torch.nn.Sequential(
+            model_seq,
+            new_head
+        )
+        
+        return joined_model, None
+
     else: 
         raise ValueError("Invalid model name")
         
@@ -120,7 +150,7 @@ def analyze_dna(username, password, sequence, model_name):
                     logits = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
                 return logits
 
-            elif model_name == 'CNN-8k-context' or model_name == 'CNN-16k-context':
+            elif 'CNN' in model_name:
                 
                 SEQUENCE_LENGTH = 8000 if '8k' in model_name else (16000 if '16k' in model_name else None)
                 pad_char = 'N'
@@ -191,8 +221,8 @@ demo = gr.Interface(
         gr.Textbox(label="DNA Sequence"),
         gr.Dropdown(label="Model", choices=[
             "GENA-Bert",
-            "CNN-8k-context",
-            "CNN-16k-context"
+            "CNN-m2-8k-context",
+            "CNN-m4-16k-context"
         ])
     ],
     outputs=[