Update app.py

app.py

@@ -20,7 +20,6 @@ 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())
@@ -186,22 +185,6 @@ def load_model(model, path="gpt_model.pth"):
 
 load_model(model)
 
-# Generate Response
-def generate_response(model, query, max_length=200):
-    model.eval()
-    with torch.no_grad():  # Disable gradient tracking
-        src = torch.tensor(tokenizer.encode(query)).unsqueeze(0).to(device)
-        tgt = torch.tensor([[1]]).to(device)  # < SOS >
-
-        for _ in range(max_length):
-            output = model(src, tgt)
-            next_token = output[:, -1, :].argmax(dim=-1, keepdim=True)
-            tgt = torch.cat([tgt, next_token], dim=1)
-            if next_token.item() == 2:  # <EOS>
-                break
-
-    return tokenizer.decode(tgt.squeeze(0).tolist())
-
 # Translation model parameters
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 MAX_LEN = 350
@@ -243,18 +226,6 @@ 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):
@@ -318,83 +289,79 @@ 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 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)
-    
-    # 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(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}")
-    # Create dummy vocabularies
-    en_vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2, '<unk>': 3, 'hello': 4, 'world': 5}
-    te_vocab = {'<pad>': 0, '<sos>': 1, '<eos>': 2, '<unk>': 3, 'హలో': 4, 'ప్రపంచం': 5}
-    te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
-    translation_available = False
-
-# 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
-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
-        
+# Initialize vocabularies from model checkpoint
+translation_available = False
+telugu_model_loaded = False
+en_vocab = None
+te_vocab = None
+te_inv_vocab = None
+model_telugu = None
+
+# Load translation model and extract vocabularies
+model_path = "english_telugu_transformer.pth"
+if os.path.exists(model_path):
     try:
-        checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
+        print("Loading Telugu translation model...")
+        checkpoint = torch.load(model_path, map_location='cpu')
         
-        # Check if vocabulary sizes match
+        # Extract vocabulary sizes from the saved model
         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)
+            # Create model with correct vocabulary sizes
+            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()
+            
+            # Try to load translation data to build 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)
+                
+                print("Building vocabularies from data...")
+                en_vocab = build_vocab(df_telugu['response'], tokenize_en, MIN_FREQ)
+                te_vocab = build_vocab(df_telugu['translated_response'], tokenize_te, MIN_FREQ)
+                te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
+                
+                # Check if vocabulary sizes match
+                if len(en_vocab) == saved_en_vocab_size and len(te_vocab) == saved_te_vocab_size:
+                    translation_available = True
+                    telugu_model_loaded = True
+                    print(f"Telugu translation model loaded successfully")
+                    print(f"English vocab size: {len(en_vocab)}, Telugu vocab size: {len(te_vocab)}")
+                else:
+                    print(f"Vocabulary size mismatch - Data EN: {len(en_vocab)}, TE: {len(te_vocab)}")
+                    print("Creating placeholder vocabularies...")
+                    # Create vocabularies with correct sizes
+                    en_vocab = {f'word_{i}': i for i in range(saved_en_vocab_size)}
+                    te_vocab = {f'word_{i}': i for i in range(saved_te_vocab_size)}
+                    te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
+                    translation_available = True
+                    telugu_model_loaded = True
+                    
+            except Exception as e:
+                print(f"Error loading Telugu dataset: {e}")
+                print("Creating placeholder vocabularies...")
+                # Create placeholder vocabularies with correct sizes
+                en_vocab = {f'word_{i}': i for i in range(saved_en_vocab_size)}
+                te_vocab = {f'word_{i}': i for i in range(saved_te_vocab_size)}
+                te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
+                translation_available = True
+                telugu_model_loaded = True
         
-        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
+        translation_available = False
+        telugu_model_loaded = False
+else:
+    print("Telugu model file not found!")
 
 # Flask App
 app = Flask(__name__)
@@ -405,20 +372,9 @@ def home():
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     return jsonify({"message": f"Welcome to TRAVIS API, Time : {current_time}"})
 
-@app.route("/intent")
-def intents():
-    try:
-        if 'intent' in df.columns:
-            unique_intents = list(set(df['intent'].dropna()))
-        else:
-            unique_intents = ["general"]  # fallback
-        return jsonify({"intents": unique_intents})
-    except Exception as e:
-        return jsonify({"error": str(e), "intents": ["general"]}), 500
-
 @app.route("/translate", methods=["POST"])
 def translate_text():
-    if not translation_available:
+    if not translation_available or not telugu_model_loaded:
         return jsonify({"error": "Translation service not available"}), 503
         
     data = request.get_json()
@@ -546,7 +502,7 @@ def query_model():
                         yield f"data: {json.dumps(response_data)}\n\n"
                 
                 # Translate to Telugu if available
-                if translation_available:
+                if translation_available and telugu_model_loaded:
                     english_response = " ".join(english_words)
                     telugu_response = translate(model_telugu, english_response, en_vocab, te_vocab, te_inv_vocab)