Update app.py

app.py

@@ -20,6 +20,7 @@ import os
 import time
 import json
 import io
+import pickle
 
 # Set PyTorch to use all available CPU threads
 torch.set_num_threads(os.cpu_count())
@@ -98,17 +99,35 @@ test_dataset = TextDataset(test_data, tokenizer)
 train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
 test_loader = DataLoader(test_dataset, batch_size=8)
 
-# Improved GPT-Style Transformer Model with optimizations
+# Function to detect model architecture from saved file
+def detect_model_architecture(model_path):
+    try:
+        checkpoint = torch.load(model_path, map_location='cpu')
+        # Check the feedforward dimension from the transformer layers
+        for key in checkpoint.keys():
+            if 'transformer.layers.0.linear1.weight' in key:
+                feedforward_dim = checkpoint[key].shape[0]  # Output dimension of first linear layer
+                embed_size = checkpoint[key].shape[1]      # Input dimension (embed_size)
+                return embed_size, feedforward_dim
+        return 256, 1024  # Default values
+    except Exception as e:
+        print(f"Could not detect architecture: {e}")
+        return 256, 1024
+
+# Improved GPT-Style Transformer Model with configurable architecture
 class GPTModel(nn.Module):
-    def __init__(self, vocab_size, embed_size=256, num_heads=8, num_layers=6, max_len=200):
+    def __init__(self, vocab_size, embed_size=256, num_heads=8, num_layers=6, max_len=200, feedforward_dim=None):
         super(GPTModel, self).__init__()
+        if feedforward_dim is None:
+            feedforward_dim = embed_size * 4
+            
         self.embedding = nn.Embedding(vocab_size, embed_size)
         self.pos_embedding = nn.Parameter(torch.randn(1, max_len, embed_size))
         self.transformer = nn.TransformerDecoder(
             nn.TransformerDecoderLayer(
                 d_model=embed_size, 
                 nhead=num_heads,
-                dim_feedforward=embed_size * 4,  # Increased feedforward dimension
+                dim_feedforward=feedforward_dim,  # Use detected or provided feedforward dimension
                 dropout=0.1,
                 batch_first=True  # Enable batch first for better performance
             ), 
@@ -134,27 +153,36 @@ class GPTModel(nn.Module):
         output = self.transformer(tgt_emb, src_emb, tgt_mask=tgt_mask)
         return self.fc_out(output)
 
-# Initialize Model with optimizations
+# Initialize Model with proper architecture detection
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-model = GPTModel(tokenizer.vocab_size).to(device)
+# Detect architecture from saved model
+model_path = "gpt_model.pth"
+if os.path.exists(model_path):
+    embed_size, feedforward_dim = detect_model_architecture(model_path)
+    print(f"Detected model architecture: embed_size={embed_size}, feedforward_dim={feedforward_dim}")
+    model = GPTModel(tokenizer.vocab_size, embed_size=embed_size, feedforward_dim=feedforward_dim).to(device)
+else:
+    model = GPTModel(tokenizer.vocab_size).to(device)
+
 optimizer = optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0.01)  # Added weight decay
 criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
 
-# Remove JIT compilation as it can cause issues with dynamic models
-# model = torch.jit.script(model)  # Commented out
-
 def load_model(model, path="gpt_model.pth"):
     if os.path.exists(path):
         try:
-            model.load_state_dict(torch.load(path, map_location=device))
+            checkpoint = torch.load(path, map_location=device)
+            model.load_state_dict(checkpoint)
             model.eval()
-            print("Model loaded successfully.")
+            print("GPT Model loaded successfully.")
+            return True
         except Exception as e:
-            print(f"Error loading model: {e}")
+            print(f"Error loading GPT model: {e}")
+            return False
     else:
-        print("Model file not found!")
+        print("GPT Model file not found!")
+        return False
 
 load_model(model)
 
@@ -215,6 +243,18 @@ def build_vocab(sentences, tokenizer, min_freq):
             vocab[word] = len(vocab)
     return vocab
 
+# Save and load vocabulary functions
+def save_vocab(vocab, path):
+    with open(path, 'wb') as f:
+        pickle.dump(vocab, f)
+
+def load_vocab(path):
+    try:
+        with open(path, 'rb') as f:
+            return pickle.load(f)
+    except:
+        return None
+
 # ==== Dataset ====
 class TranslationDataset(Dataset):
     def __init__(self, df, en_vocab, te_vocab):
@@ -278,19 +318,32 @@ def translate(model, sentence, en_vocab, te_vocab, te_inv_vocab, max_len=MAX_LEN
     translated = [te_inv_vocab[idx.item()] for idx in tgt_ids[0][1:]]
     return ' '.join(translated[:-1]) if translated[-1] == '<eos>' else ' '.join(translated)
 
-# ==== Load Translation Data ====
+# ==== Load Translation Data and Vocabularies ====
 try:
     df_telugu = pd.read_csv("merged_translated_responses.csv")
     df_telugu = df_telugu.dropna(subset=['response', 'translated_response'])
     df_telugu['response'] = df_telugu['response'].astype(str)
     df_telugu['translated_response'] = df_telugu['translated_response'].astype(str)
     
-    # Build vocabularies
-    en_vocab = build_vocab(df_telugu['response'], tokenize_en, MIN_FREQ)
-    te_vocab = build_vocab(df_telugu['translated_response'], tokenize_te, MIN_FREQ)
+    # Try to load saved vocabularies first
+    en_vocab = load_vocab('en_vocab.pkl')
+    te_vocab = load_vocab('te_vocab.pkl')
+    
+    if en_vocab is None or te_vocab is None:
+        print("Building new vocabularies...")
+        # Build vocabularies
+        en_vocab = build_vocab(df_telugu['response'], tokenize_en, MIN_FREQ)
+        te_vocab = build_vocab(df_telugu['translated_response'], tokenize_te, MIN_FREQ)
+        # Save vocabularies
+        save_vocab(en_vocab, 'en_vocab.pkl')
+        save_vocab(te_vocab, 'te_vocab.pkl')
+    else:
+        print("Loaded saved vocabularies")
+    
     te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
     
-    print("Telugu translation dataset loaded successfully")
+    print(f"Telugu translation dataset loaded successfully")
+    print(f"English vocab size: {len(en_vocab)}, Telugu vocab size: {len(te_vocab)}")
     translation_available = True
 except Exception as e:
     print(f"Error loading Telugu dataset: {e}")
@@ -300,17 +353,47 @@ except Exception as e:
     te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
     translation_available = False
 
-# Initialize Translation Model
+# Initialize Translation Model with correct vocabulary sizes
 model_telugu = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
                                  len(en_vocab), len(te_vocab), NHEAD, FFN_HID_DIM).to(DEVICE)
 
 # Load saved weights for translation model
-try:
-    model_telugu.load_state_dict(torch.load("english_telugu_transformer.pth", map_location=torch.device('cpu')))
-    model_telugu.eval()
-    print("Telugu translation model loaded successfully")
-except Exception as e:
-    print(f"Error loading Telugu translation model: {e}")
+def load_telugu_model():
+    model_path = "english_telugu_transformer.pth"
+    if not os.path.exists(model_path):
+        print("Telugu model file not found!")
+        return False
+        
+    try:
+        checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+        
+        # Check if vocabulary sizes match
+        if 'src_tok_emb.weight' in checkpoint:
+            saved_en_vocab_size = checkpoint['src_tok_emb.weight'].shape[0]
+            saved_te_vocab_size = checkpoint['tgt_tok_emb.weight'].shape[0]
+            current_en_vocab_size = len(en_vocab)
+            current_te_vocab_size = len(te_vocab)
+            
+            print(f"Saved model vocabs - EN: {saved_en_vocab_size}, TE: {saved_te_vocab_size}")
+            print(f"Current model vocabs - EN: {current_en_vocab_size}, TE: {current_te_vocab_size}")
+            
+            if saved_en_vocab_size != current_en_vocab_size or saved_te_vocab_size != current_te_vocab_size:
+                print("Vocabulary size mismatch! Creating new model with saved vocabulary sizes...")
+                global model_telugu
+                model_telugu = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
+                                                saved_en_vocab_size, saved_te_vocab_size, NHEAD, FFN_HID_DIM).to(DEVICE)
+        
+        model_telugu.load_state_dict(checkpoint)
+        model_telugu.eval()
+        print("Telugu translation model loaded successfully")
+        return True
+    except Exception as e:
+        print(f"Error loading Telugu translation model: {e}")
+        return False
+
+# Load Telugu model
+telugu_model_loaded = load_telugu_model()
+if not telugu_model_loaded:
     translation_available = False
 
 # Flask App
@@ -378,6 +461,7 @@ def translate_text():
 def generate_text():
     data = request.get_json()
     query = data.get("query", "")
+    print("entered /generate")
     if not query:
         return jsonify({"error": "Query cannot be empty"}), 400
 
@@ -513,4 +597,5 @@ def get_audio():
 if __name__ == "__main__":
     print("Starting Flask application...")
     print(f"Translation service available: {translation_available}")
+    print(f"Telugu model loaded: {telugu_model_loaded}")
     app.run(host="0.0.0.0", debug=True)