276 lines
9.1 KiB
Python
276 lines
9.1 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 pydantic import BaseModel
|
|
import os
|
|
|
|
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")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model architecture
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class RotaryPositionalEmbedding(nn.Module):
|
|
def __init__(self, dim=256, max_seq_len=512):
|
|
super().__init__()
|
|
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
|
|
self.register_buffer("inv_freq", inv_freq)
|
|
self.max_seq_len = max_seq_len
|
|
|
|
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)
|
|
emb = torch.cat((freqs, freqs), dim=-1)
|
|
return emb[: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 _apply_rope(self, x, freqs):
|
|
return x * freqs.cos() + self._rotate_half(x) * freqs.sin()
|
|
|
|
def _rotate_half(self, x):
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
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, max_seq_len=512):
|
|
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()
|
|
for _ in range(max_new_tokens):
|
|
logits = self(input_ids)
|
|
logits = logits[:, -1, :]
|
|
|
|
if temperature > 0:
|
|
logits = logits / temperature
|
|
|
|
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)
|
|
input_ids = torch.cat([input_ids, next_token], dim=-1)
|
|
|
|
if eos_token_id is not None and next_token.item() == eos_token_id:
|
|
break
|
|
|
|
return input_ids
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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['<EOS>']
|
|
)
|
|
|
|
generated_text = vocab.indices_to_text(generated[0].tolist())
|
|
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
|
|
|
|
|
|
@app.on_event("startup")
|
|
def load_artifacts():
|
|
global model, vocab
|
|
if not os.path.exists(MODEL_PATH):
|
|
print(f"[WARN] {MODEL_PATH} not found. Model loading skipped.")
|
|
return
|
|
if not os.path.exists(VOCAB_PATH):
|
|
print(f"[WARN] {VOCAB_PATH} not found. Vocab loading skipped.")
|
|
return
|
|
|
|
try:
|
|
import pickle
|
|
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=True)
|
|
vocab = pickle.load(open(VOCAB_PATH, "rb"))
|
|
|
|
model = GPTModelWithRoPE(
|
|
vocab_size=len(vocab.word2idx) if hasattr(vocab, 'word2idx') else 42962
|
|
).to(DEVICE)
|
|
|
|
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
else:
|
|
model.load_state_dict(checkpoint)
|
|
|
|
model.eval()
|
|
print(f"[OK] Model loaded on {DEVICE}")
|
|
print(f"[OK] Vocab size: {len(vocab.word2idx)}")
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to load artifacts: {e}")
|
|
model = None
|
|
vocab = None
|
|
|
|
|
|
@app.get("/api/health")
|
|
def health():
|
|
return {
|
|
"status": "ok",
|
|
"model_loaded": model is not None,
|
|
"vocab_loaded": vocab is not None,
|
|
"device": str(DEVICE),
|
|
}
|
|
|
|
|
|
@app.post("/api/generate", response_model=GenerateResponse)
|
|
def generate(req: GenerateRequest):
|
|
if model is None or vocab is None:
|
|
raise HTTPException(status_code=503, detail="Model or vocab not loaded")
|
|
|
|
texts = []
|
|
for _ in range(req.num_samples):
|
|
text = generate_clickbait(
|
|
model, vocab,
|
|
prompt=req.prompt,
|
|
temperature=req.temperature,
|
|
top_k=req.top_k,
|
|
)
|
|
texts.append(text)
|
|
|
|
return GenerateResponse(texts=texts)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entry point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
port = int(os.environ.get("PORT", 8000))
|
|
uvicorn.run(app, host="0.0.0.0", port=port)
|