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))