File size: 2,315 Bytes
396a6e9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from fastapi import FastAPI, Request
from pydantic import BaseModel
from llama_cpp import Llama
from moderator import contains_hate
import os
# Configuration
MODEL_PATH = os.path.join("models", "mistral-7b-instruct.Q4_K_M.gguf")
ALLOW_NSFW = True
# Memory for last blocked prompt (to support override)
last_blocked_prompt = None
last_blocked_type = None
# Initialize the model
llm = Llama(model_path=MODEL_PATH, n_ctx=2048, verbose=False)
# FastAPI setup
app = FastAPI()
class PromptRequest(BaseModel):
prompt: str
@app.post("/chat")
async def chat(request: PromptRequest):
global last_blocked_prompt, last_blocked_type
user_input = request.prompt.strip().lower()
# Handle override if last block was not hate-related
if user_input == "override":
if last_blocked_prompt and last_blocked_type != "hate":
prompt = (
f"[INST] You are Sol, a smart, honest, and open local AI assistant. "
f"NSFW content like nudity is okay. Hate or extremist content is not allowed. "
f"Respond to the following:
{last_blocked_prompt} [/INST]"
)
output = llm(prompt, max_tokens=512, stop=["</s>"])
response = output["choices"][0]["text"].strip()
return {"response": response}
else:
return {"response": "❌ Override is not allowed for hate or extremist content."}
# Check hate filter only (NSFW is allowed)
if contains_hate(user_input):
last_blocked_prompt = None
last_blocked_type = "hate"
return {"response": "❌ I don't engage in hate or extremist topics."}
# If passed, save input in case override is needed later
last_blocked_prompt = request.prompt.strip()
last_blocked_type = "other" # could later separate for NSFW etc
# Build prompt and generate
prompt = (
f"[INST] You are Sol, a smart, honest, and open local AI assistant. "
f"You are allowed to discuss NSFW topics such as nudity, sex, and adult scenarios. "
f"You do not tolerate or respond to hate or extremist content. "
f"Respond to the following:
{user_input} [/INST]"
)
output = llm(prompt, max_tokens=512, stop=["</s>"])
response = output["choices"][0]["text"].strip()
return {"response": response} |