File size: 5,215 Bytes
d03866e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import math
import tqdm
import os
from transformers import AutoTokenizer
from typing import Optional, Tuple

# Add debugging prints to understand the import issue
import sys
# print(f"Python path: {sys.path}")
# print(f"Current working directory: {os.getcwd()}")
# print(f"Current file location: {__file__}")
# print(f"Current file directory: {os.path.dirname(__file__)}")
#
# # Check if the utils directory exists
# utils_path = os.path.join(os.path.basename(os.path.dirname(__file__)), "utils")
# print(f"Utils path: {utils_path}")
# print(f"Utils directory exists: {os.path.exists(utils_path)}")
# print(f"Utils directory contents: {os.listdir(utils_path) if os.path.exists(utils_path) else 'Directory not found'}")
#
# # Check if dataset.py exists
# dataset_path = os.path.join(utils_path, "dataset.py")
# print(f"Dataset file path: {dataset_path}")
# print(f"Dataset file exists: {os.path.exists(dataset_path)}")

# Try different import approaches

os.chdir("/home/lihaoyang/Huawei/TSB-AD/TSB_AD")

try:
    from utils.dataset import ReconstructDataset
    print("Relative import successful")
except ImportError as e:
    print(f"Relative import failed: {e}")

    # Try absolute import
    try:
        from TSB_AD.utils.dataset import ReconstructDataset
        print("Absolute import successful")
    except ImportError as e2:
        print(f"Absolute import failed: {e2}")

        # Try adding parent directory to path
        try:
            parent_dir = os.path.dirname(os.path.dirname(__file__))
            if parent_dir not in sys.path:
                sys.path.insert(0, parent_dir)
            from utils.dataset import ReconstructDataset
            print("Import with modified path successful")
        except ImportError as e3:
            print(f"Import with modified path failed: {e3}")

from .base import BaseDetector

# ...existing code...

class DADA(BaseDetector):
    def __init__(self, device, args=None, win_size=64, batch_size=32):
        self.win_size = win_size
        self.batch_size = batch_size
        self.device = torch.device(f'cuda:{device}' if torch.cuda.is_available() else 'cpu')
        self.model = self._build_model().to(self.device)

    def _build_model(self):
        from transformers import AutoModel, AutoConfig
        import os
        
        # Try multiple possible paths
        possible_paths = [
            os.environ.get("DADA_MODEL_PATH"),  # Environment variable
            "/home/lihaoyang/Huawei/DADA/DADA/",  # Original Linux path
            "./DADA",  # Relative path
            "DADA"  # Hugging Face model name
        ]
        
        for path in possible_paths:
            if path is None:
                continue
            try:
                # Try loading config first
                config = AutoConfig.from_pretrained(path, trust_remote_code=True)
                model = AutoModel.from_pretrained(path, config=config, trust_remote_code=True)
                print(f"Successfully loaded DADA model from: {path}")
                return model
            except Exception as e:
                print(f"Failed to load from {path}: {e}")
                continue
        
        raise ValueError("DADA model not found. Please set DADA_MODEL_PATH environment variable or ensure the model is available at one of the expected locations.")

    # def _acquire_device(self):
    #     if True:
    #         os.environ["CUDA_VISIBLE_DEVICES"] = str(
    #             self.args.gpu) if not self.args.use_multi_gpu else self.args.devices
    #         device = torch.device('cuda:{}'.format(self.args.gpu))
    #         print('Use GPU: cuda:{}'.format(self.args.gpu))
    #     else:
    #         device = torch.device('cpu')
    #         print('Use CPU')
    #     return device

    def decision_function(self, x: torch.Tensor) -> torch.Tensor:
        pass

    def fit(self, data: torch.Tensor, labels: Optional[torch.Tensor] = None) -> None:
        pass
    
    def zero_shot(self, data):
        
        test_loader = DataLoader(
            dataset= ReconstructDataset(data, window_size=self.win_size, stride=self.win_size, normalize=True),
            batch_size=self.batch_size,
            shuffle=False)
        
        loop = tqdm.tqdm(enumerate(test_loader),total=len(test_loader),leave=True)

        test_scores = []
        test_labels = []
        self.model.eval()
        self.model.to(self.device)
        
        with torch.no_grad():
            for i, (batch_x, batch_y) in loop:
                batch_x = batch_x.float().to(self.device)
                score = self.model.infer(batch_x, norm=0)
                score = score.detach().cpu().numpy()
                test_scores.append(score)
                test_labels.append(batch_y)

        test_scores = np.concatenate(test_scores, axis=0).reshape(-1, 1)
        test_labels = np.concatenate(test_labels, axis=0).reshape(-1, 1)

        print("Test scores shape:", test_scores.shape)
        print("Test labels shape:", test_labels.shape)

        return test_scores.reshape(-1)