Dauren777
871147b688
- server.py: FastAPI serves model (best_generator.pt) + static files - /api/generate endpoint with temperature/top_k controls - /api/health endpoint for status - Frontend auto-detects backend, shows connection status - API URL configurable from UI (persisted in localStorage) - NaN/inf safety in generation - Clean output (strips BOS/EOS/PAD/UNK tokens)
359 lines
12 KiB
Python
359 lines
12 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel
|
|
import os, pickle, sys
|
|
|
|
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
MODEL_PATH = os.environ.get("MODEL_PATH", "best_generator.pt")
|
|
VOCAB_PATH = os.environ.get("VOCAB_PATH", "vocab.pt")
|
|
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "7"))
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model architecture
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class RotaryPositionalEmbedding(nn.Module):
|
|
def __init__(self, dim=32):
|
|
super().__init__()
|
|
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim).float() / dim))
|
|
self.register_buffer("inv_freq", inv_freq)
|
|
|
|
def forward(self, x):
|
|
seq_len = x.shape[1]
|
|
t = torch.arange(seq_len, device=x.device).float()
|
|
freqs = torch.einsum("i,j->ij", t, self.inv_freq)
|
|
return freqs[:seq_len]
|
|
|
|
|
|
class MultiHeadAttentionWithRoPE(nn.Module):
|
|
def __init__(self, d_model=256, n_heads=8, dropout=0.1):
|
|
super().__init__()
|
|
assert d_model % n_heads == 0
|
|
self.d_model = d_model
|
|
self.n_heads = n_heads
|
|
self.d_head = d_model // n_heads
|
|
|
|
self.W_q = nn.Linear(d_model, d_model, bias=True)
|
|
self.W_k = nn.Linear(d_model, d_model, bias=True)
|
|
self.W_v = nn.Linear(d_model, d_model, bias=True)
|
|
self.W_o = nn.Linear(d_model, d_model, bias=True)
|
|
|
|
self.rope = RotaryPositionalEmbedding(dim=self.d_head)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
def _rotate_half(self, x):
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
def _apply_rope(self, x, freqs):
|
|
return x * freqs.cos() + self._rotate_half(x) * freqs.sin()
|
|
|
|
def forward(self, x, mask=None):
|
|
B, T, C = x.shape
|
|
Q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
|
|
K = self.W_k(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
|
|
V = self.W_v(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
|
|
|
|
freqs = self.rope(x)
|
|
Q = self._apply_rope(Q, freqs)
|
|
K = self._apply_rope(K, freqs)
|
|
|
|
attn = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5)
|
|
if mask is not None:
|
|
attn = attn.masked_fill(mask == 0, float("-inf"))
|
|
attn = F.softmax(attn, dim=-1)
|
|
attn = self.dropout(attn)
|
|
|
|
out = (attn @ V).transpose(1, 2).contiguous().view(B, T, C)
|
|
return self.W_o(out)
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, d_model=256, d_ff=1024, dropout=0.1):
|
|
super().__init__()
|
|
self.linear1 = nn.Linear(d_model, d_ff, bias=True)
|
|
self.linear2 = nn.Linear(d_ff, d_model, bias=True)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
def forward(self, x):
|
|
return self.linear2(self.dropout(F.gelu(self.linear1(x))))
|
|
|
|
|
|
class DecoderBlockWithRoPE(nn.Module):
|
|
def __init__(self, d_model=256, n_heads=8, d_ff=1024, dropout=0.1):
|
|
super().__init__()
|
|
self.self_attention = MultiHeadAttentionWithRoPE(d_model, n_heads, dropout)
|
|
self.feed_forward = FeedForward(d_model, d_ff, dropout)
|
|
self.norm1 = nn.LayerNorm(d_model)
|
|
self.norm2 = nn.LayerNorm(d_model)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
def forward(self, x, mask=None):
|
|
x = x + self.dropout(self.self_attention(self.norm1(x), mask))
|
|
x = x + self.dropout(self.feed_forward(self.norm2(x)))
|
|
return x
|
|
|
|
|
|
class GPTModelWithRoPE(nn.Module):
|
|
def __init__(self, vocab_size=42962, d_model=256, n_heads=8,
|
|
n_layers=3, d_ff=1024, dropout=0.1):
|
|
super().__init__()
|
|
self.token_embedding = nn.Embedding(vocab_size, d_model)
|
|
self.embedding_dropout = nn.Dropout(dropout)
|
|
self.decoder_blocks = nn.ModuleList([
|
|
DecoderBlockWithRoPE(d_model, n_heads, d_ff, dropout)
|
|
for _ in range(n_layers)
|
|
])
|
|
self.final_norm = nn.LayerNorm(d_model)
|
|
self.output_projection = nn.Linear(d_model, vocab_size, bias=True)
|
|
|
|
def forward(self, x, mask=None):
|
|
x = self.embedding_dropout(self.token_embedding(x))
|
|
for block in self.decoder_blocks:
|
|
x = block(x, mask)
|
|
x = self.final_norm(x)
|
|
logits = self.output_projection(x)
|
|
return logits
|
|
|
|
@torch.no_grad()
|
|
def generate(self, input_ids, max_new_tokens=7, temperature=1.0,
|
|
top_k=50, top_p=0.9, eos_token_id=None):
|
|
self.eval()
|
|
B, T = input_ids.shape
|
|
# Causal mask for current sequence length
|
|
mask = torch.tril(torch.ones(1, 1, T, T, device=input_ids.device))
|
|
|
|
for _ in range(max_new_tokens):
|
|
logits = self(input_ids, mask)
|
|
logits = logits[:, -1, :]
|
|
|
|
if temperature > 0:
|
|
logits = logits / temperature
|
|
|
|
# Replace nan/inf with -inf to avoid crashes
|
|
logits = torch.nan_to_num(logits, nan=float("-inf"), posinf=float("-inf"), neginf=float("-inf"))
|
|
|
|
if top_k > 0:
|
|
values, _ = torch.topk(logits, top_k, dim=-1)
|
|
logits[logits < values[:, -1:]] = float("-inf")
|
|
|
|
if top_p < 1.0:
|
|
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
|
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
|
sorted_indices_to_remove = cumulative_probs > top_p
|
|
sorted_logits[sorted_indices_to_remove] = float("-inf")
|
|
logits = sorted_logits.scatter(1, sorted_indices, sorted_logits)
|
|
|
|
probs = F.softmax(logits, dim=-1)
|
|
next_token = torch.multinomial(probs, num_samples=1)
|
|
|
|
# Update causal mask for new token
|
|
input_ids = torch.cat([input_ids, next_token], dim=-1)
|
|
new_T = input_ids.shape[1]
|
|
mask = torch.tril(torch.ones(1, 1, new_T, new_T, device=input_ids.device))
|
|
|
|
if eos_token_id is not None and next_token.item() == eos_token_id:
|
|
break
|
|
|
|
return input_ids
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Vocab (in case no vocab file exists)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class DummyVocab:
|
|
"""Minimal vocab that passes through token IDs as text."""
|
|
word2idx = {'<BOS>': 0, '<EOS>': 1, '<PAD>': 2, '<UNK>': 3}
|
|
idx2word = {0: '<BOS>', 1: '<EOS>', 2: '<PAD>', 3: '<UNK>'}
|
|
|
|
def __init__(self):
|
|
for i in range(4, 42962):
|
|
self.word2idx[f'<TOK_{i}>'] = i
|
|
self.idx2word[i] = f'<TOK_{i}>'
|
|
|
|
def text_to_indices_for_generator(self, text):
|
|
return [self.word2idx.get(w, 3) for w in text.strip().split()]
|
|
|
|
def indices_to_text(self, indices):
|
|
return ' '.join(self.idx2word.get(i, '<UNK>') for i in indices)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Generation function
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@torch.no_grad()
|
|
def generate_clickbait(model, vocab, prompt="", max_new_tokens=7,
|
|
temperature=1.0, top_k=50, top_p=0.9):
|
|
model.eval()
|
|
|
|
if prompt:
|
|
tokens = vocab.text_to_indices_for_generator(prompt)
|
|
input_ids = torch.LongTensor([[vocab.word2idx['<BOS>']] + tokens]).to(DEVICE)
|
|
else:
|
|
input_ids = torch.LongTensor([[vocab.word2idx['<BOS>']]]).to(DEVICE)
|
|
|
|
generated = model.generate(
|
|
input_ids,
|
|
max_new_tokens=max_new_tokens,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
eos_token_id=vocab.word2idx.get('<EOS>', 1)
|
|
)
|
|
|
|
generated_text = vocab.indices_to_text(generated[0].tolist())
|
|
|
|
# Clean up special tokens for display
|
|
for tok in ['<BOS>', '<EOS>', '<PAD>', '<UNK>']:
|
|
generated_text = generated_text.replace(tok, '')
|
|
generated_text = generated_text.strip()
|
|
|
|
return generated_text
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
app = FastAPI(title="ClickBait Generator API")
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
class GenerateRequest(BaseModel):
|
|
prompt: str = ""
|
|
temperature: float = 1.0
|
|
top_k: int = 50
|
|
num_samples: int = 1
|
|
|
|
|
|
class GenerateResponse(BaseModel):
|
|
texts: list[str]
|
|
|
|
|
|
model = None
|
|
vocab = None
|
|
using_dummy_vocab = False
|
|
|
|
|
|
@app.on_event("startup")
|
|
def load_artifacts():
|
|
global model, vocab, using_dummy_vocab
|
|
if not os.path.exists(MODEL_PATH):
|
|
print(f"[WARN] {MODEL_PATH} not found. Model loading skipped.")
|
|
return
|
|
|
|
try:
|
|
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=True)
|
|
|
|
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
|
|
state_dict = checkpoint['model_state_dict']
|
|
else:
|
|
state_dict = checkpoint
|
|
|
|
vocab_size = state_dict['token_embedding.weight'].shape[0]
|
|
model = GPTModelWithRoPE(vocab_size=vocab_size).to(DEVICE)
|
|
|
|
# Remove causal_mask (we compute it dynamically)
|
|
sd_clean = {k: v for k, v in state_dict.items() if k != 'causal_mask'}
|
|
missing, unexpected = model.load_state_dict(sd_clean, strict=False)
|
|
if missing:
|
|
print(f"[WARN] Missing keys: {missing}")
|
|
if unexpected:
|
|
print(f"[WARN] Unexpected keys: {unexpected}")
|
|
|
|
model.eval()
|
|
|
|
# Load vocab
|
|
if os.path.exists(VOCAB_PATH):
|
|
vocab = pickle.load(open(VOCAB_PATH, "rb"))
|
|
using_dummy_vocab = False
|
|
print(f"[OK] Vocab loaded: {len(vocab.word2idx)} tokens")
|
|
else:
|
|
print(f"[WARN] {VOCAB_PATH} not found, using fallback dummy vocab")
|
|
vocab = DummyVocab()
|
|
using_dummy_vocab = True
|
|
|
|
print(f"[OK] Model loaded on {DEVICE} | params: {sum(p.numel() for p in model.parameters()):,}")
|
|
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to load model: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
|
|
@app.get("/api/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"model_loaded": model is not None,
|
|
"vocab_loaded": vocab is not None,
|
|
"using_dummy_vocab": using_dummy_vocab,
|
|
"device": str(DEVICE),
|
|
"vocab_size": len(vocab.word2idx) if vocab else 0,
|
|
"note": "vocab.pt not found — output shows token IDs" if using_dummy_vocab else "ready",
|
|
}
|
|
|
|
|
|
@app.post("/api/generate", response_model=GenerateResponse)
|
|
def generate(req: GenerateRequest):
|
|
if model is None:
|
|
raise HTTPException(status_code=503, detail="Model not loaded")
|
|
|
|
# Clamp / validate
|
|
temperature = max(0.1, min(3.0, req.temperature))
|
|
top_k = max(1, min(200, req.top_k))
|
|
count = max(1, min(20, req.num_samples))
|
|
|
|
texts = []
|
|
for _ in range(count):
|
|
text = generate_clickbait(
|
|
model, vocab or DummyVocab(),
|
|
prompt=req.prompt,
|
|
temperature=temperature,
|
|
top_k=top_k,
|
|
max_new_tokens=MAX_NEW_TOKENS,
|
|
)
|
|
texts.append(text)
|
|
|
|
return GenerateResponse(texts=texts)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Static files + SPA fallback
|
|
# ---------------------------------------------------------------------------
|
|
|
|
STATIC_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
@app.get("/")
|
|
def serve_index():
|
|
return FileResponse(os.path.join(STATIC_DIR, "index.html"))
|
|
|
|
@app.get("/{path:path}")
|
|
def serve_static(path: str):
|
|
file_path = os.path.join(STATIC_DIR, path)
|
|
if os.path.isfile(file_path):
|
|
return FileResponse(file_path)
|
|
return FileResponse(os.path.join(STATIC_DIR, "index.html"))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
port = int(os.environ.get("PORT", 8000))
|
|
uvicorn.run(app, host="0.0.0.0", port=port)
|