mlai / app.py
saifeddinemk's picture
Fixed app v2
bc119db
raw
history blame
1.21 kB
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
from typing import List
# Initialize the FastAPI app
app = FastAPI()
# 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):
role: str
content: str
class MessagesInput(BaseModel):
messages: List[Message]
@app.post("/generate/")
async def generate_response(messages_input: MessagesInput):
try:
# 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 from the response
generated_text = response[0]["generated_text"]
return {
"response": generated_text
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# Run the app
# To start the server, use the command: uvicorn filename:app --host 0.0.0.0 --port 8000