from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import GPT2LMHeadModel, AutoTokenizer
import torch, uuid

device = "cuda" if torch.cuda.is_available() else "cpu"
tok    = AutoTokenizer.from_pretrained("model", use_fast=False)
model  = GPT2LMHeadModel.from_pretrained("model").to(device).eval()

app = FastAPI()
memory: dict[str, list[tuple[str, str]]] = {}

class Msg(BaseModel):
    text: str

@torch.inference_mode()
@app.post("/chat")
def chat(req: Request, m: Msg):

    sid = req.headers.get("X-Session-ID")            
    if not sid:                                      
        sid = str(uuid.uuid4())
    history = memory.get(sid, [])[-5:]               

    system_prompt = (
        "あなたは「oupe ec」"
        "齋藤凪沙が書いた日記や思考、発話のスタイルを学習しており、"
        "彼女の語感・比喩・距離感・揺らぎをもとに、質問や発話に返答します。"
        "意識やキャラ性は演じません。"
        "あくまで、齋藤凪沙という人が考えたら出しそうな言葉の構造を予測して返します。\\n----\\n"
    )

    buf = [system_prompt]
    for u, o in history:
        buf.append(f"<USER> {u}\\n<OPE>  {o}")
    buf.append(f"<USER> {m.text.strip()}\\n<OPE>")
    prompt = "\\n".join(buf)

    ids = tok(prompt, return_tensors="pt").to(device)
    prompt_len = ids.input_ids.shape[-1]

    out = model.generate(
        **ids,
        max_new_tokens=96,
        do_sample=True,
        temperature=0.9,
        top_p=0.90,
        repetition_penalty=1.15,
        pad_token_id=tok.eos_token_id,
    )

    gen_tokens = out[0][prompt_len:]
    reply = tok.decode(gen_tokens, skip_special_tokens=True).strip()
    reply = reply.strip('\\"“”「」『』') 

    memory.setdefault(sid, []).append((m.text, reply))

    return {"reply": reply, "session_id": sid}

正直記録用にいじってみたけど、うまく動かなくて履歴の保存できなかったから、また改良しないといけない。自動保存はとりあえず後回しでもいいかな。