Update app.py

app.py

@@ -80,7 +80,7 @@ def load_model(model_name: str):
             model_seq,
             new_head
         )
-        weights = torch.load('/CNN_1stGEAC_m2_best.pth')
+        weights = torch.load('CNN_1stGEAC_m2_best.pth')
         model.load_state_dict(weights)
         return model, None
 
@@ -120,9 +120,9 @@ 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':
+            elif model_name == 'CNN-8k-context' or model_name == 'CNN-16k-context':
                 
-                SEQUENCE_LENGTH = 8000
+                SEQUENCE_LENGTH = 8000 if '8k' in model_name else (16000 if '16k' in model_name else None)
                 pad_char = 'N'
                 
                 # Truncate sequence
@@ -190,8 +190,9 @@ demo = gr.Interface(
         gr.Textbox(label="Password", type="password"),
         gr.Textbox(label="DNA Sequence"),
         gr.Dropdown(label="Model", choices=[
-            "gena-bert",
-            "CNN"
+            "GENA-Bert",
+            "CNN-8k-context",
+            "CNN-16k-context"
         ])
     ],
     outputs=[