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['']] + tokens]).to(DEVICE) else: input_ids = torch.LongTensor([[vocab.word2idx['']]]).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[''] ) 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)