Spaces:
Running
Running
File size: 3,234 Bytes
11463c3 | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | import json
import time
from pathlib import Path
import torch
from model import (
GreesyGPT,
generate_moderation,
ReasoningMode,
OutputFormat,
DEVICE,
)
# βββββββββββββββββββββββββββββββββββββββββββββ
# Model Initialization
# βββββββββββββββββββββββββββββββββββββββββββββ
model = GreesyGPT()
weights_path = Path(__file__).parent / "greesy_gpt.pt"
if weights_path.exists():
model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
print(f"Loaded weights from {weights_path}")
else:
print("No trained weights found, using fresh initialization.")
model.to(DEVICE)
model.eval()
# βββββββββββββββββββββββββββββββββββββββββββββ
# OpenAIβstyle Chat Completion Wrapper
# βββββββββββββββββββββββββββββββββββββββββββββ
def chat_completions(
model: GreesyGPT,
messages,
reasoning_mode: ReasoningMode = ReasoningMode.LOW,
):
"""
Emulates the OpenAI Chat Completions API format.
Input:
messages = [
{"role": "system", "content": "..."},
{"role": "user", "content": "..."}
]
Output:
{
"id": "...",
"object": "chat.completion",
"choices": [...]
}
"""
user_message = ""
system_message = ""
for m in messages:
if m["role"] == "user":
user_message = m["content"]
elif m["role"] == "system":
system_message = m["content"]
result = generate_moderation(
model,
prompt=user_message,
mode=reasoning_mode,
output_format=OutputFormat.MARKDOWN,
)
verdict = result["verdict"]
thinking = result["thinking"]
completion_text = verdict
response = {
"id": f"greesy-{int(time.time()*1000)}",
"object": "chat.completion",
"created": int(time.time()),
"model": "latest",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": completion_text,
},
"finish_reason": "stop",
}
],
"metadata": {
"reasoning": thinking,
},
}
return response
# βββββββββββββββββββββββββββββββββββββββββββββ
# Example Usage
# βββββββββββββββββββββββββββββββββββββββββββββ
if __name__ == "__main__":
messages = [
{"role": "system", "content": "You are a moderation system."},
{"role": "user", "content": "You're so stupid, nobody likes you."},
]
response = chat_completions(
model,
messages,
reasoning_mode=ReasoningMode.MEDIUM,
)
print(json.dumps(response, indent=2)) |