coder-1/Small_Language_Model_Of_Banking

1import gradio as gr2import torch3import torch.nn as nn4import torch.nn.functional as func5import math6from dataclasses import dataclass7import tiktoken8 9 10SECRET_KEY = "Shreeramchandraji-GOD-1000BCE"11device = 'cuda' if torch.cuda.is_available() else 'cpu'12class LayerNorm(nn.Module):13    def __init__(self, ndimension, bias):14        super().__init__()15        self.weight = nn.Parameter(torch.ones(ndimension))16        self.bias = nn.Parameter(torch.zeros(ndimension)) if bias else None17    18    def forward(self, x):19        return func.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)20 21class Attention(nn.Module):22    def __init__(self, config):23        super().__init__()24        assert config.n_embd % config.n_head == 025        self.c_attn = nn.Linear(config.n_embd, 3*config.n_embd, bias=config.bias)26        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)27        self.attn_dropout = nn.Dropout(config.dropout)28        self.resid_dropout = nn.Dropout(config.dropout)29        self.n_head = config.n_head30        self.n_embd = config.n_embd31        self.flash = hasattr(func, 'scaled_dot_product_attention')32        if not self.flash:33            self.register_buffer('bias', torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))34    35    def forward(self, x):36        B, T, C = x.size()37        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)38        k = k.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)39        q = q.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)40        v = v.view(B, T, self.n_head, C//self.n_head).transpose(1, 2)41        if self.flash:42            y = func.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.attn_dropout.p if self.training else 0.0, is_causal=True)43        else:44            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))45            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))46            att = func.softmax(att, dim=-1)47            att = self.attn_dropout(att)48            y = att @ v49        y = y.transpose(1, 2).contiguous().view(B, T, C)50        y = self.resid_dropout(self.c_proj(y))51        return y52 53class MLP(nn.Module):54    def __init__(self, config):55        super().__init__()56        self.c_fc = nn.Linear(config.n_embd, 4*config.n_embd, bias=config.bias)57        self.gelu = nn.GELU()58        self.c_proj = nn.Linear(4*config.n_embd, config.n_embd, bias=config.bias)59        self.dropout = nn.Dropout(config.dropout)60    61    def forward(self, x):62        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))63 64class Block(nn.Module):65    def __init__(self, config):66        super().__init__()67        self.ln1 = LayerNorm(config.n_embd, config.bias)68        self.attn = Attention(config)69        self.ln2 = LayerNorm(config.n_embd, config.bias)70        self.mlp = MLP(config)71    72    def forward(self, x):73        x = x + self.attn(self.ln1(x))74        x = x + self.mlp(self.ln2(x))75        return x76 77@dataclass78class Setting_Up_GPT:79    vocab_size: int80    block_size: int81    n_layer: int82    n_head: int83    n_embd: int84    dropout: float = 0.085    bias: bool = True86 87class GPT(nn.Module):88    def __init__(self, config):89        super().__init__()90        self.config = config91        self.transformer = nn.ModuleDict(dict(92            wte=nn.Embedding(config.vocab_size, config.n_embd),93            wpe=nn.Embedding(config.block_size, config.n_embd),94            drop=nn.Dropout(config.dropout),95            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),96            ln_f=LayerNorm(config.n_embd, config.bias)97        ))98        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)99        self.transformer.wte.weight = self.lm_head.weight100        self.apply(self._init_weights)101        for pn, p in self.named_parameters():102            if pn.endswith('c_proj.weight'):103                nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2*config.n_layer))104    105    def _init_weights(self, module):106        if isinstance(module, nn.Linear):107            nn.init.normal_(module.weight, mean=0.0, std=0.02)108            if module.bias is not None:109                nn.init.zeros_(module.bias)110        elif isinstance(module, nn.Embedding):111            nn.init.normal_(module.weight, mean=0.0, std=0.02)112    113    def forward(self, idx, targets=None):114        device = idx.device115        b, t = idx.size()116        assert t <= self.config.block_size117        pos = torch.arange(0, t, dtype=torch.long, device=device)118        tok_emb = self.transformer.wte(idx)119        pos_emb = self.transformer.wpe(pos)120        x = self.transformer.drop(tok_emb + pos_emb)121        for bl in self.transformer.h:122            x = bl(x)123        x = self.transformer.ln_f(x)124        if targets is not None:125            logits = self.lm_head(x)126            loss = func.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)127            return logits, loss128        else:129            logits = self.lm_head(x[:, [-1], :])130            return logits, None131 132config = Setting_Up_GPT(133    vocab_size=50257,134    block_size=128,135    n_layer=6,136    n_head=8,137    n_embd=512,138    dropout=0.1,139    bias=True140)141 142model = GPT(config)143 144# Loading your trained weights145# Upload 'best_model_params.pt' to your HuggingFace Space146try:147    model.load_state_dict(torch.load('best_model_params.pt', map_location=device, weights_only=True))148    # model.load_state_dict(torch.load('best_model_params.pt', map_location=device))149    print("✅ Model weights loaded successfully!")150except:151    print("⚠️ Model weights not found. Upload 'best_model_params.pt' to your Space.")152 153model.to(device)154model.eval()155 156# Load tokenizer157convertor = tiktoken.get_encoding('gpt2')158 159# ===== GENERATION FUNCTION =====160 161def generate_answer(query, max_tokens=150, temperature=0.7):162    """Generate answer for a given query"""163    prompt = f"Query: {query}\nAnswer:"164    input_ids = convertor.encode_ordinary(prompt)165    input_tensor = torch.tensor([input_ids], dtype=torch.long).to(device)166    generated_ids = input_ids.copy()167    168    with torch.no_grad():169        for _ in range(max_tokens):170            logits, _ = model(input_tensor)171            next_token_logits = logits[0, -1, :] / temperature172            probs = torch.nn.functional.softmax(next_token_logits, dim=-1)173            next_token = torch.multinomial(probs, num_samples=1).item()174            generated_ids.append(next_token)175            176            decoded = convertor.decode(generated_ids)177            if '<|endoftext|>' in decoded or '\nQuery:' in decoded:178                break179            180            input_tensor = torch.cat([181                input_tensor,182                torch.tensor([[next_token]], dtype=torch.long).to(device)183            ], dim=1)184    185    full_output = convertor.decode(generated_ids)186    187    if 'Answer:' in full_output:188        answer = full_output.split('Answer:')[1]189        answer = answer.split('<|endoftext|>')[0]190        answer = answer.split('\nQuery:')[0]191        answer = answer.strip()192    else:193        answer = "I couldn't generate a response."194    195    return answer196 197# ===== GRADIO API WRAPPER =====198 199def api_generate(query, api_key, max_tokens=150, temperature=0.7):200    """Gradio wrapper with API key authentication"""201    if not query.strip():202        return "❌ Error: Please enter a query"203    204    if api_key != SECRET_KEY:205        return "❌ Error: Invalid API key"206    207    try:208        answer = generate_answer(query, int(max_tokens), temperature)209        return answer210    except Exception as e:211        return f"❌ Error: {str(e)}"212 213# ===== GRADIO INTERFACE =====214 215demo = gr.Interface(216    fn=api_generate,217    inputs=[218        gr.Textbox(219            label="💬 Query", 220            placeholder="e.g., How do I activate my credit card?",221            lines=3222        ),223        gr.Textbox(224            label="🔑 API Key", 225            type="password",226            placeholder="Enter your API key"227        ),228        gr.Slider(229            minimum=50, 230            maximum=300, 231            value=150, 232            step=10,233            label="📏 Max Tokens"234        ),235        gr.Slider(236            minimum=0.1,237            maximum=2.0,238            value=0.7,239            step=0.1,240            label="🌡️ Temperature"241        )242    ],243    outputs=gr.Textbox(244        label="🎯 AI Response",245        lines=5246    ),247    title="🎙️ Banking SLM API for Voice Agents",248    description="Small Language Model trained on banking queries. Designed for voice agent integration.",249    examples=[250        ["How do I activate my new credit card?", SECRET_KEY, 150, 0.7],251        ["What are your customer support hours?", SECRET_KEY, 100, 0.7],252    ],253    api_name="generate"254)255 256if __name__ == "__main__":257    demo.launch()