app.py

This file contains code to

app.py

@@ -0,0 +1,86 @@
+import torch
+from torch import nn
+from tokenizers import Tokenizer
+import torch.nn.functional as F
+import gradio as gr
+
+# Enhanced Model with Custom Attention
+class RachanaLLM(nn.Module):
+    def __init__(self, vocab_size=50000, embed_dim=284, num_heads=4, num_layers=4, dropout=0.2, max_len=256):
+        super(RachanaLLM, self).__init__()
+        self.embed_dim = embed_dim
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.positional_encoding = nn.Parameter(torch.zeros(1, max_len, embed_dim))
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True)
+        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.output_layer = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, src, src_mask=None):
+        src = self.embedding(src) * torch.sqrt(torch.tensor(self.embed_dim, dtype=torch.float32))
+        src = src + self.positional_encoding[:, :src.size(1)]
+        output = self.transformer_encoder(src, src_key_padding_mask=src_mask)
+        return self.output_layer(output)
+
+# Load model and tokenizer
+checkpoint = torch.load("best_model.pth")  # Use your actual path
+model = RachanaLLM()
+model.load_state_dict(checkpoint['model_state'] if 'model_state' in checkpoint else checkpoint)
+model.eval()
+tokenizer = Tokenizer.from_file("telugu_tokenizer_50k.json")  # Use your actual path
+
+# Advanced Beam Search with Temperature, Top-K Sampling, and Dynamic Repetition Penalty
+def beam_search_decoder(model, tokenizer, input_ids, beam_width=3, max_length=20, temperature=1.0, top_k=5, repetition_penalty=1.2):
+    model.eval()
+    with torch.no_grad():
+        sequences = [[input_ids, 0.0]]
+        
+        for _ in range(max_length):
+            all_candidates = []
+            for seq, score in sequences:
+                outputs = model(seq)
+                logits = outputs[:, -1] / temperature
+
+                # If a token has been used already, penalize its logit
+                for token_id in set(seq[0].tolist()):
+                    if token_id in seq[0]:
+                        logits[0][token_id] /= repetition_penalty
+                
+                top_logits, top_indices = torch.topk(logits, top_k, dim=-1)
+                softmax_scores = F.log_softmax(top_logits, dim=-1)
+
+                for i in range(top_k):
+                    next_token_id = top_indices[0][i].item()
+                    if seq[0].tolist().count(next_token_id) < 2: # This ensures a token is penalized only if it appears more than once.
+                        candidate = [torch.cat([seq, top_indices[:, i:i+1]], dim=1), score + softmax_scores[0][i].item()]
+                        all_candidates.append(candidate)
+                    
+            # Sort candidates by score
+            ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True)
+            sequences = ordered[:beam_width]
+        
+        best_seq = sequences[0][0]
+        result_text = tokenizer.decode(best_seq[0].tolist())
+        return result_text
+
+
+
+# Predict function for Gradio Interface
+def predict(input_sentence):
+    input_ids = torch.tensor(tokenizer.encode(input_sentence).ids).unsqueeze(0)
+    generated_text = beam_search_decoder(model, tokenizer, input_ids, temperature=0.8, top_k=10)
+    return generated_text
+
+# Gradio Interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=2, placeholder="Type a sentence..."),  # Updated line
+    outputs="text",
+    examples=[
+        "ఈ రోజు వాతావరణం చాలా బాగుంది.",
+        "సినిమా బాగుందా లేదా చెప్పు!",
+        "ఆలయ అధికారులు దర్శన ఏర్పాట్లు చేశారు."
+    ],
+    title="Rachana LLM"
+)
+
+iface.launch(share=True)