Update app.py

app.py

@@ -1,29 +1,25 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from src.model_inference import predict
+from src.utils import setup_logging, log_info, log_error
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("canstralian/CyberAttackDetection")
-model = AutoModelForCausalLM.from_pretrained("canstralian/CyberAttackDetection")
+# Set up logging
+setup_logging()
 
 # Define the input data model
 class LogData(BaseModel):
     log: str
 
 @app.post("/predict")
-async def predict(data: LogData):
-    # Tokenize the input log data
-    inputs = tokenizer(data.log, return_tensors="pt")
-    
-    # Generate predictions
-    with torch.no_grad():
-        outputs = model.generate(**inputs)
-    
-    # Decode the generated tokens to text
-    prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return {"prediction": prediction}
+async def predict_route(data: LogData):
+    try:
+        # Perform prediction
+        prediction = predict(data.log)
+        log_info(f'Prediction: {prediction}')
+        return {"prediction": prediction}
+    except Exception as e:
+        log_error(f'An error occurred: {e}')
+        return {"error": str(e)}