model_tensorflow

1name: model_tensorflow
2version: 1.0.0
3runtime: tensorflow
4description: TensorFlow Custom Model
5parameters:
6  layers:
7    type: integer
8    default: 3
9    description: Number of dense layers
10  units:
11    type: integer
12    default: 64
13    description: Units per layer
14

1
2import pandas as pd
3import numpy as np
4import tensorflow as tf
5from typing import Dict, Any
6
7class Model:
8    def __init__(self):
9        self.model = None
10        self.input_dim = 10 
11
12    def _build_model(self, input_dim):
13        """
14        Build a simple TF Autoencoder
15        """
16        model = tf.keras.Sequential([
17            tf.keras.layers.Dense(16, activation='relu', input_shape=(input_dim,)),
18            tf.keras.layers.Dense(8, activation='relu'),
19            tf.keras.layers.Dense(16, activation='relu'),
20            tf.keras.layers.Dense(input_dim, activation='linear')
21        ])
22        model.compile(optimizer='adam', loss='mse')
23        return model
24
25    def train(self, ctx) -> Dict[str, Any]:
26        """
27        Train TensorFlow model
28        """
29        ctx.logger.info("Starting TensorFlow training...")
30        
31        if ctx.df is None or ctx.df.empty:
32            ctx.logger.warning("No data, generating dummy")
33            X = np.random.randn(100, 10).astype(np.float32)
34        else:
35            X = ctx.df.select_dtypes(include=[np.number]).values.astype(np.float32)
36            
37        self.input_dim = X.shape[1]
38        self.model = self._build_model(self.input_dim)
39        
40        history = self.model.fit(X, X, epochs=10, batch_size=32, verbose=0)
41        final_loss = history.history['loss'][-1]
42        
43        ctx.logger.info(f"Training completed. Loss: {final_loss}")
44        
45        # self.model.save("tf_model") # In real app, save to artifacts
46        
47        return {
48            "status": "success",
49            "model_type": "TensorFlow Autoencoder",
50            "final_loss": float(final_loss)
51        }
52
53    def infer(self, ctx) -> pd.DataFrame:
54        """
55        Inference
56        """
57        ctx.logger.info("Starting TensorFlow inference...")
58        
59        if ctx.df is None or ctx.df.empty:
60            X = np.random.randn(20, self.input_dim).astype(np.float32)
61            ids = [f"user_{i}" for i in range(20)]
62        else:
63            X = ctx.df.select_dtypes(include=[np.number]).values.astype(np.float32)
64             # Handle dimension mismatch
65            if X.shape[1] != self.input_dim:
66                 if X.shape[1] > self.input_dim:
67                     X = X[:, :self.input_dim]
68                 else:
69                     padding = np.zeros((X.shape[0], self.input_dim - X.shape[1]), dtype=np.float32)
70                     X = np.hstack((X, padding))
71            
72            if "entity_id" in ctx.df.columns:
73                ids = ctx.df["entity_id"].values
74            else:
75                ids = [f"entity_{i}" for i in range(len(X))]
76                
77        if self.model is None:
78             self.model = self._build_model(self.input_dim)
79
80        reconstructions = self.model.predict(X, verbose=0)
81        mse = np.mean(np.power(X - reconstructions, 2), axis=1)
82        
83        results = []
84        for i, score in enumerate(mse):
85            risk = min(100.0, float(score) * 50)
86            results.append({
87                "entity_id": str(ids[i]),
88                "risk_score": float(risk),
89                "anomaly_type": "tf_reconstruction_error" if risk > 50 else "normal",
90                "details": {"mse": float(score)}
91            })
92            
93        return pd.DataFrame(results)
94
95    def execute(self, data=None):
96         # shim for v1
97        class MockCtx:
98            def __init__(self, d): self.df = d if d else pd.DataFrame(); self.logger = type('obj', (object,), {'info': print, 'warning': print})
99        return self.infer(MockCtx(pd.DataFrame(data) if data else None)).to_dict('records')
100