Fixed app v2

app.py

@@ -1,16 +1,22 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import pipeline
+from llama_cpp import Llama
+from functools import lru_cache
 import uvicorn
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Load the text generation pipeline with the specified model
-try:
-    pipe = pipeline("text-generation", model="ammarnasr/codegen2-1B-security", trust_remote_code=True)
-except Exception as e:
-    raise RuntimeError(f"Failed to load model: {e}")
+# Lazy load the Llama model
+@lru_cache(maxsize=1)
+def load_model():
+    try:
+        return Llama.from_pretrained(
+            repo_id="prithivMLmods/Llama-3.2-1B-GGUF",
+            filename="Llama-3.2-1B.F16.gguf"
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model: {e}")
 
 # Define request model for log data
 class LogRequest(BaseModel):
@@ -23,6 +29,7 @@ class AnalysisResponse(BaseModel):
 # Define the route for security log analysis
 @app.post("/analyze_security_logs", response_model=AnalysisResponse)
 async def analyze_security_logs(request: LogRequest):
+    llm = load_model()
     try:
         # Security-focused prompt
         prompt = (
@@ -32,11 +39,18 @@ async def analyze_security_logs(request: LogRequest):
             f"{request.log_data}"
         )
 
-        # Generate response from the pipeline with a controlled max length
-        response = pipe(prompt, max_length=512, num_return_sequences=1)
+        # Generate response from the Llama model
+        response = llm.create_chat_completion(
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
+        )
 
         # Extract and return the analysis text
-        analysis_text = response[0]["generated_text"]
+        analysis_text = response["choices"][0]["message"]["content"]
         return AnalysisResponse(analysis=analysis_text)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))