Fixed app v2

app.py

@@ -1,79 +1,43 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer
+from transformers import pipeline
 import uvicorn
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Configure and load the quantized model
-model_id = 'bakch92/Llama-3.1-security'
+# 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}")
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-)
-
-# Load tokenizer and model with 4-bit quantization settings
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-tokenizer.pad_token = tokenizer.eos_token
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    quantization_config=bnb_config,
-    device_map="auto",
-)
-model.eval()
-
-# Define request and response models
-class SecurityLogRequest(BaseModel):
+# Define request model for log data
+class LogRequest(BaseModel):
     log_data: str
 
-class SecurityAnalysisResponse(BaseModel):
+# Define response model
+class AnalysisResponse(BaseModel):
     analysis: str
 
-# Inference function
-def generate_response(input_text: str) -> str:
-    streamer = TextStreamer(tokenizer=tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
-    messages = [
-        {"role": "system", "content": "You are an information security AI assistant specialized in analyzing security logs. Identify potential threats, suspicious IP addresses, unauthorized access attempts, and recommend actions based on the logs."},
-        {"role": "user", "content": f"Please analyze the following security logs and provide insights on any potential malicious activity:\n{input_text}"}
-    ]
-    
-    input_ids = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_tensors="pt",
-    ).to(model.device)
-    
-    # Generate response with the model
-    outputs = model.generate(
-        input_ids,
-        streamer=streamer,
-        max_new_tokens=512,  # Limit max tokens for faster response
-        num_beams=1,
-        do_sample=True,
-        temperature=0.1,
-        top_p=0.95,
-        top_k=10
-    )
-    
-    # Extract and return generated text
-    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response_text
-
 # Define the route for security log analysis
-@app.post("/analyze_security_logs", response_model=SecurityAnalysisResponse)
-async def analyze_security_logs(request: SecurityLogRequest):
+@app.post("/analyze_security_logs", response_model=AnalysisResponse)
+async def analyze_security_logs(request: LogRequest):
     try:
-        # Run inference
-        analysis_text = generate_response(request.log_data)
-        return SecurityAnalysisResponse(analysis=analysis_text)
+        # Security-focused prompt
+        prompt = (
+            "Analyze the following network log data for any indicators of malicious activity, "
+            "such as unusual IP addresses, unauthorized access attempts, data exfiltration, or anomalies. "
+            "Provide details on potential threats, IPs involved, and suggest actions if any threats are detected.\n\n"
+            f"{request.log_data}"
+        )
+
+        # Generate response from the pipeline with a controlled max length
+        response = pipe(prompt, max_length=512, num_return_sequences=1)
+
+        # Extract and return the analysis text
+        analysis_text = response[0]["generated_text"]
+        return AnalysisResponse(analysis=analysis_text)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))