Fixed app v2

app.py

@@ -1,13 +1,13 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import pipeline
-from typing import List, Dict
+from typing import List
 
 # Initialize the FastAPI app
 app = FastAPI()
 
-# Initialize the text generation pipeline
-pipe = pipeline("text-generation", model="CyberNative-AI/Colibri_8b_v0.1")
+# Initialize the text generation pipeline with the specified model
+pipe = pipeline("text-generation", model="jcordon5/Mistral-7B-cybersecurity-rules")
 
 # Define the input schema for FastAPI
 class Message(BaseModel):
@@ -20,13 +20,13 @@ class MessagesInput(BaseModel):
 @app.post("/generate/")
 async def generate_response(messages_input: MessagesInput):
     try:
-        # Convert messages to the expected format
+        # Convert the messages to the expected format for the pipeline
         messages = [{"role": msg.role, "content": msg.content} for msg in messages_input.messages]
 
         # Generate response using the pipeline
         response = pipe(messages)
         
-        # Extract generated text
+        # Extract generated text from the response
         generated_text = response[0]["generated_text"]
 
         return {