Update app.py

app.py

@@ -1,11 +1,15 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForSeq2SeqLM,
+)
 import torch
 
 # Define the model names and mappings
 MODEL_MAPPING = {
-    "text2shellcommands": "Canstralian/text2shellcommands",
-    "pentest_ai": "Canstralian/pentest_ai",
+    "text2shellcommands": "t5-small",  # Example seq2seq model
+    "pentest_ai": "bert-base-uncased",  # Example sequence classification model
 }
 
 # Sidebar for model selection
@@ -18,13 +22,9 @@ def select_model():
 @st.cache_resource
 def load_model_and_tokenizer(model_name):
     try:
-        # Use a fallback model for testing
-        if model_name == "Canstralian/text2shellcommands":
-            model_name = "t5-small"
-
         # Load the tokenizer and model
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if "seq2seq" in model_name.lower():
+        if "t5" in model_name or "seq2seq" in model_name:
             model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
         else:
             model = AutoModelForSequenceClassification.from_pretrained(model_name)
@@ -38,7 +38,7 @@ def load_model_and_tokenizer(model_name):
 # Handle predictions
 def predict_with_model(user_input, model, tokenizer, model_choice):
     if model_choice == "text2shellcommands":
-        # Generate shell commands
+        # Generate shell commands (seq2seq task)
         inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
         with torch.no_grad():
             outputs = model.generate(**inputs)
@@ -79,4 +79,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()