Fixed app v2

app.py

@@ -1,70 +1,86 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import List, Dict, Any
+from typing import List
 from transformers import pipeline
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the custom CyBERTuned model directly from Hugging Face
+
 unmasker = pipeline("fill-mask", model="s2w-ai/CyBERTuned-SecurityLLM")
 
-# Define request model for multiple log entries
+# Define request model
 class LogRequest(BaseModel):
-    logs: List[str]  # Expecting a list of log entries
+    log: str
 
 # Define response model
 class ThreatResponse(BaseModel):
     log: str
     prompt: str
-    pred: Dict[str, List[Dict[str, Any]]]  # Dictionary with structured lists for predictions
+    #threat_level_predictions: List[str]
+    #threat_type_predictions: List[str]
+    #detected_threat_level: str
+    #detected_threat_type: str
+    pred : List[object]
 
-# Function to predict masked words for threat level, type, SRC IP, DEST IP, and Protocol
-def predict_threat(log: str, unmasker, topk=5) -> Dict[str, List[Dict[str, Any]]]:
-    # Create prompt with masked tokens for each category
-    prompt = (
-        f"{log} Threat level: <mask>. Threat type: <mask>. "
-        f"Attack type: <mask>. SRC IP: <mask>. DEST IP: <mask>. Protocol: <mask>."
-    )
+# Function to predict masked words for threat level and type
+def predict_threat(log: str, unmasker, topk=5) -> List[List[str]]:
+    # Create prompt with masked tokens for threat level and threat type
+    prompt = f"{log}"
     
     # Predict top options for each <mask>
     predictions = unmasker(prompt, top_k=topk)
     
-    # Separate predictions for each masked category
-    threat_level_predictions = predictions[:topk]
-    threat_type_predictions = predictions[topk:2*topk]
-    attack_type_predictions = predictions[2*topk:3*topk]
-    src_ip_predictions = predictions[3*topk:4*topk]
-    dest_ip_predictions = predictions[4*topk:5*topk]
-    protocol_predictions = predictions[5*topk:6*topk]
+    # Extract top predictions for each <mask>
+    #threat_level_predictions = [pred["token_str"].strip() for pred in predictions[:topk]]
+    #threat_type_predictions = [pred["token_str"].strip() for pred in predictions[topk:2*topk]]
+    
+    return predictions
+def get_maximum_predictions(data):
+    # Initialize list to store maximum values for each prediction array
+    max_predictions = []
+
+    # Loop over each prediction array in "pred"
+    for index, predictions in enumerate(data["pred"]):
+        max_score = float('-inf')
+        max_prediction = None
+
+        # Find the prediction with the highest score in the current array
+        for pred in predictions:
+            if pred["score"] > max_score:
+                max_score = pred["score"]
+                max_prediction = pred["token_str"].strip()
 
-    # Structure the response in a dictionary, with each entry containing a list of dictionaries
-    return {
-        "threat_level_predictions": [pred for pred in threat_level_predictions],
-        "threat_type_predictions": [pred for pred in threat_type_predictions],
-        "attack_type_predictions": [pred for pred in attack_type_predictions],
-        "src_ip_predictions": [pred for pred in src_ip_predictions],
-        "dest_ip_predictions": [pred for pred in dest_ip_predictions],
-        "protocol_predictions": [pred for pred in protocol_predictions]
-    }
+        # Append the result with the max prediction for this array
+        max_predictions.append({
+            f"max_prediction_{index + 1}": max_prediction
+        })
 
-# FastAPI endpoint for detecting threat level and type for multiple logs
-@app.post("/detect_threat", response_model=List[ThreatResponse])
+    return max_predictions
+
+# Get result
+ 
+# FastAPI endpoint for detecting threat level and type
+@app.post("/detect_threat", response_model=ThreatResponse)
 async def detect_threat(log_request: LogRequest):
-    responses = []
-    for log in log_request.logs:
-        # Predict the threat level and type for each log entry
-        predictions = predict_threat(log, unmasker)
+    log = log_request.log
+
+    # Predict the threat level and type for the given log entry
+    predictions = predict_threat(log, unmasker)
+    
+    # Extract top predictions for threat level and type
+    ##threat_level_predictions = predictions[0] if len(predictions) > 0 else ["Unknown"]
+   ## threat_type_predictions = predictions[1] if len(predictions) > 1 else ["Unknown"]
 
-        # Prepare response for each log entry
-        response = ThreatResponse(
-            log=log,
-            prompt=(
-                f"{log} Threat level: <mask>. Threat type: <mask>. "
-                f"Attack type: <mask>. SRC IP: <mask>. DEST IP: <mask>. Protocol: <mask>."
-            ),
-            pred=predictions
-        )
-        responses.append(response)
+    # Use the top prediction as the most likely threat level and type
+    ##detected_threat_level = threat_level_predictions[0] if threat_level_predictions else "Unknown"
+    #detected_threat_type = threat_type_predictions[0] if threat_type_predictions else "Unknown"
+
+    # Prepare response
+    response = ThreatResponse(
+        log=log,
+        prompt=f"{log} Threat level: <mask>. Threat type: <mask>.",
+        pred=get_maximum_predictions(predictions)
+    )
 
-    return responses
+    return response