Fixed app v2

app.py

@@ -6,7 +6,7 @@ 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
@@ -17,20 +17,25 @@ class LogRequest(BaseModel):
 class ThreatResponse(BaseModel):
     log: str
     prompt: str
-    pred: List[str]
+    #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 and type
-def predict_threat(log: str, unmasker, topk=5) -> List[str]:
+def predict_threat(log: str, unmasker, topk=5) -> List[List[object]]:
     # Create prompt with masked tokens for threat level and threat type
     prompt = f"{log} Threat level most likely is: <mask>. Most likely the Threat type: <mask>."
     
     # Predict top options for each <mask>
     predictions = unmasker(prompt, top_k=topk)
     
-    # Extract only the token_str values to return a list of strings
-    pred_strings = [pred["token_str"].strip() for pred in predictions]
+    # 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 pred_strings
+    return predictions
 
 # FastAPI endpoint for detecting threat level and type
 @app.post("/detect_threat", response_model=ThreatResponse)
@@ -38,13 +43,21 @@ async def detect_threat(log_request: LogRequest):
     log = log_request.log
 
     # Predict the threat level and type for the given log entry
-    pred_strings = predict_threat(log, unmasker)
+    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"]
+
+    # 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=pred_strings
+        pred=predictions
     )
 
     return response