Spaces:
Running
Running
File size: 3,112 Bytes
5cda0af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
import unittest
from unittest.mock import patch, MagicMock
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import streamlit as st
import io
class TestStreamlitApp(unittest.TestCase):
@patch("transformers.AutoTokenizer.from_pretrained")
@patch("transformers.AutoModelForSequenceClassification.from_pretrained")
def test_load_model_success(self, mock_model, mock_tokenizer):
# Mock the tokenizer and model loading
mock_tokenizer.return_value = MagicMock(spec=AutoTokenizer)
mock_model.return_value = MagicMock(spec=AutoModelForSequenceClassification)
tokenizer, model = load_model("Canstralian/CyberAttackDetection")
# Assert that the tokenizer and model are not None
self.assertIsNotNone(tokenizer)
self.assertIsNotNone(model)
mock_tokenizer.assert_called_once_with("Canstralian/CyberAttackDetection")
mock_model.assert_called_once_with("Canstralian/CyberAttackDetection")
@patch("transformers.AutoTokenizer.from_pretrained")
@patch("transformers.AutoModelForSequenceClassification.from_pretrained")
def test_predict_classification(self, mock_model, mock_tokenizer):
# Mock the tokenizer and model for inference
mock_tokenizer.return_value = MagicMock(spec=AutoTokenizer)
mock_model.return_value = MagicMock(spec=AutoModelForSequenceClassification)
# Simulate model outputs
mock_model.return_value.__call__.return_value = MagicMock(logits=torch.tensor([[1.0, 2.0, 3.0]]))
# Call the prediction function
inputs = mock_tokenizer("Test input", return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
outputs = mock_model.return_value(**inputs)
logits = outputs.logits
predicted_class = torch.argmax(logits, dim=-1).item()
# Assert that the predicted class is correct
self.assertEqual(predicted_class, 2) # The class with the highest score (index 2)
@patch("transformers.AutoTokenizer.from_pretrained")
@patch("transformers.AutoModelForSeq2SeqLM.from_pretrained")
def test_generate_shell_command(self, mock_model, mock_tokenizer):
# Mock the tokenizer and model for shell command generation
mock_tokenizer.return_value = MagicMock(spec=AutoTokenizer)
mock_model.return_value = MagicMock(spec=AutoModelForSeq2SeqLM)
# Simulate model output (fake shell command)
mock_model.return_value.generate.return_value = torch.tensor([[1, 2, 3, 4]])
# Simulate text input
user_input = "Create a directory"
inputs = mock_tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
outputs = mock_model.return_value.generate(**inputs)
generated_command = mock_tokenizer.decode(outputs[0], skip_special_tokens=True)
# Assert the generated command is as expected
self.assertEqual(generated_command, "mkdir directory") # Assuming the model generates this
if __name__ == "__main__":
unittest.main()
|