Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import textwrap
 import torch
 from datetime import datetime
@@ -10,29 +9,47 @@ import pandas as pd
 from torch.utils.data import Dataset, DataLoader
 from torch.nn.utils.rnn import pad_sequence
 from sklearn.model_selection import train_test_split
-from flask import Flask ,request, jsonify,send_file,after_this_request
+from flask import Flask, request, jsonify, send_file, after_this_request, Response, stream_with_context
 from collections import Counter
 from flask_cors import CORS
 import requests
 from gtts import gTTS
-from googletrans import Translator
+
 import uuid
 import os
 import time
+import json
+import io
+
+# Set PyTorch to use all available CPU threads
+torch.set_num_threads(os.cpu_count())
+torch.set_num_interop_threads(os.cpu_count())
 
+# Enable PyTorch JIT for better performance
+torch.jit.enable_onednn_fusion(True)
 
 # Load Dataset
-df = pd.read_csv("https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY")
-df = df.dropna(subset=['instruction', 'response'])
+try:
+    df = pd.read_csv("https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY")
+    df = df.dropna(subset=['instruction', 'response'])
+    # Ensure all entries are strings
+    df['instruction'] = df['instruction'].astype(str)
+    df['response'] = df['response'].astype(str)
+    print("Main dataset loaded successfully")
+except Exception as e:
+    print(f"Error loading main dataset: {e}")
+    # Create a dummy dataset for testing
+    df = pd.DataFrame({
+        'instruction': ['Hello', 'How are you?'],
+        'response': ['Hi there!', 'I am doing well, thank you!'],
+        'intent': ['greeting', 'greeting']
+    })
 
-# Ensure all entries are strings
-df['instruction'] = df['instruction'].astype(str)
-df['response'] = df['response'].astype(str)
 # Tokenizer (Scratch)
 class ScratchTokenizer:
     def __init__(self):
-        self.word2idx = {"<PAD>": 0, "<SOS>": 1, "<EOS>": 2, "<UNK>": 3}
-        self.idx2word = {0: "<PAD>", 1: "<SOS>", 2: "<EOS>", 3: "<UNK>"}
+        self.word2idx = {"<PAD>": 0, "< SOS >": 1, "<EOS>": 2, "<UNK>": 3}
+        self.idx2word = {0: "<PAD>", 1: "< SOS >", 2: "<EOS>", 3: "<UNK>"}
         self.vocab_size = 4
 
     def build_vocab(self, texts):
@@ -81,40 +98,61 @@ 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
-
+# Improved GPT-Style Transformer Model with optimizations
 class GPTModel(nn.Module):
     def __init__(self, vocab_size, embed_size=256, num_heads=8, num_layers=6, max_len=200):
         super(GPTModel, self).__init__()
         self.embedding = nn.Embedding(vocab_size, embed_size)
         self.pos_embedding = nn.Parameter(torch.randn(1, max_len, embed_size))
-        # The problem was here, setting num_encoder_layers to 0
-        # makes the model try to access a non-existent layer.
-        # The solution is to remove the encoder completely.
-        self.transformer = nn.TransformerDecoder(nn.TransformerDecoderLayer(d_model=embed_size, nhead=num_heads), num_layers=num_layers)
+        self.transformer = nn.TransformerDecoder(
+            nn.TransformerDecoderLayer(
+                d_model=embed_size, 
+                nhead=num_heads,
+                dim_feedforward=embed_size * 4,  # Increased feedforward dimension
+                dropout=0.1,
+                batch_first=True  # Enable batch first for better performance
+            ), 
+            num_layers=num_layers
+        )
         self.fc_out = nn.Linear(embed_size, vocab_size)
+        
+        # Initialize weights for better training
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
     def forward(self, src, tgt):
         src_emb = self.embedding(src) + self.pos_embedding[:, :src.size(1), :]
         tgt_emb = self.embedding(tgt) + self.pos_embedding[:, :tgt.size(1), :]
-
-        # Causal Mask for Auto-Regressive Decoding
         tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)
-        output = self.transformer(tgt_emb.permute(1, 0, 2), src_emb.permute(1, 0, 2), tgt_mask=tgt_mask)
-        return self.fc_out(output.permute(1, 0, 2))
+        output = self.transformer(tgt_emb, src_emb, tgt_mask=tgt_mask)
+        return self.fc_out(output)
 
-# Initialize Model
+# Initialize Model with optimizations
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
 model = GPTModel(tokenizer.vocab_size).to(device)
-optimizer = optim.AdamW(model.parameters(), lr=2e-4)
+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):
-        model.load_state_dict(torch.load(path, map_location=device))
-        model.eval()
-        print("Model loaded successfully.")
+        try:
+            model.load_state_dict(torch.load(path, map_location=device))
+            model.eval()
+            print("Model loaded successfully.")
+        except Exception as e:
+            print(f"Error loading model: {e}")
     else:
         print("Model file not found!")
 
@@ -125,7 +163,7 @@ 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>
+        tgt = torch.tensor([[1]]).to(device)  # < SOS >
 
         for _ in range(max_length):
             output = model(src, tgt)
@@ -136,7 +174,7 @@ def generate_response(model, query, max_length=200):
 
     return tokenizer.decode(tgt.squeeze(0).tolist())
 
-
+# Translation model parameters
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 MAX_LEN = 350
 BATCH_SIZE = 8
@@ -149,9 +187,19 @@ NUM_EPOCHS = 18
 MIN_FREQ = 2
 
 # ==== Tokenizers ====
-spacy_eng = spacy.load("en_core_web_sm")
+try:
+    spacy_eng = spacy.load("en_core_web_sm")
+    print("Spacy English model loaded successfully")
+except OSError:
+    print("Warning: Spacy English model not found. Using simple tokenizer.")
+    spacy_eng = None
+
 def tokenize_en(text):
-    return [tok.text.lower() for tok in spacy_eng.tokenizer(text)]
+    if spacy_eng:
+        return [tok.text.lower() for tok in spacy_eng.tokenizer(text)]
+    else:
+        # Simple fallback tokenizer
+        return text.lower().split()
 
 def tokenize_te(text):
     return text.strip().split(" ")
@@ -189,13 +237,6 @@ class TranslationDataset(Dataset):
 
         return torch.tensor(en_ids), torch.tensor(te_ids)
 
-# ==== Collate Function ====
-def collate_fn(batch):
-    src_batch, tgt_batch = zip(*batch)
-    src_batch = pad_sequence(src_batch, padding_value=en_vocab['<pad>'], batch_first=True)
-    tgt_batch = pad_sequence(tgt_batch, padding_value=te_vocab['<pad>'], batch_first=True)
-    return src_batch, tgt_batch
-
 # ==== Transformer Model ====
 class Seq2SeqTransformer(nn.Module):
     def __init__(self, num_encoder_layers, num_decoder_layers,
@@ -237,144 +278,239 @@ 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 Data ====
-df_telugu = pd.read_csv("merged_translated_responses.csv")  # columns: 'en', 'te'
-# Clean NaN or non-string entries
-df_telugu = df_telugu.dropna(subset=['response', 'translated_response'])
-
-# Ensure all entries are strings
-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)
-te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
-
-# Prepare Dataset & DataLoader
-dataset = TranslationDataset(df_telugu, en_vocab, te_vocab)
-dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
-
-# Initialize Model
-# model = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
-#                            len(en_vocab), len(te_vocab), NHEAD, FFN_HID_DIM).to(DEVICE)
-
-pad_idx = te_vocab['<pad>']
-criterion_telugu = nn.CrossEntropyLoss(ignore_index=pad_idx)
-optimizer_telugu = optim.Adam(model.parameters(), lr=0.0005)
-
-# ==== Training ====
-# for epoch in range(NUM_EPOCHS):
-#     loss = train(model, dataloader, optimizer, criterion)
-#     print(f"Epoch {epoch+1}, Loss: {loss:.4f}")
-
-# ==== Try Translation ====
-
-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
-model_telugu.load_state_dict(torch.load("english_telugu_transformer.pth",map_location = torch.device('cpu')))
-model_telugu.eval()
-app=Flask(__name__)
+# ==== Load Translation Data ====
+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)
+    te_inv_vocab = {idx: tok for tok, idx in te_vocab.items()}
+    
+    print("Telugu translation dataset loaded successfully")
+    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
+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}")
+    translation_available = False
+
+# Flask App
+app = Flask(__name__)
 CORS(app)
 
-
 @app.route("/")
 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():
-    return jsonify({"intents" :list(set(df['intent'].dropna()))})
-
-
+    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:
+        return jsonify({"error": "Translation service not available"}), 503
+        
     data = request.get_json()
     text = data.get("text", "")
-    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    print("Entered '/translate' at time: ",current_time)
     if not text:
         return jsonify({"error": "Text cannot be empty"}), 400
 
-    # First generate English response
-    english_response = text
-    start=time.time()
-    # Then translate to Telugu
-    telugu_response = translate(model_telugu, english_response, en_vocab, te_vocab, te_inv_vocab)
-    end=time.time()
-    return jsonify({
-        "english": english_response,
-        "telugu": telugu_response,
-        "time": end-start
-    })
+    def generate():
+        try:
+            start = time.time()
+            word_count = 0
+            
+            # Translate to Telugu word by word
+            telugu_response = translate(model_telugu, text, en_vocab, te_vocab, te_inv_vocab)
+            
+            # Stream each word of the translation
+            for word in telugu_response.split():
+                word_count += 1
+                response_data = {
+                    "word": word.strip(),
+                    "timestamp": time.time() - start,
+                    "word_count": word_count,
+                    "type": "telugu"
+                }
+                yield f"data: {json.dumps(response_data)}\n\n"
+        except Exception as e:
+            error_data = {"error": str(e), "type": "error"}
+            yield f"data: {json.dumps(error_data)}\n\n"
+    
+    return Response(
+        stream_with_context(generate()),
+        mimetype='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive'
+        }
+    )
 
 @app.route("/generate", methods=["POST"])
 def generate_text():
     data = request.get_json()
     query = data.get("query", "")
-    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    print("Entered '/generate' at time: ",current_time)
-
     if not query:
         return jsonify({"error": "Query cannot be empty"}), 400
-    start=time.time()
-    response = generate_response(model, query)
-    end=time.time()
-    # Clean the response
-    def clean_response(response):
-        return response.replace("<EOS>", "").replace("<SOS>", "").strip()
-    
-    response = clean_response(response)
+
+    def generate():
+        try:
+            start = time.time()
+            word_count = 0
+            model.eval()
+            
+            with torch.no_grad():
+                src = torch.tensor(tokenizer.encode(query)).unsqueeze(0).to(device)
+                tgt = torch.tensor([[1]]).to(device)  # < SOS >
+                
+                for _ in range(200):  # 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
+                        
+                    word = tokenizer.idx2word.get(next_token.item(), "<UNK>")
+                    if word not in ["<PAD>", "<EOS>", "< SOS >"]:
+                        word_count += 1
+                        response_data = {
+                            "word": word.strip(),
+                            "timestamp": time.time() - start,
+                            "word_count": word_count,
+                            "type": "english"
+                        }
+                        yield f"data: {json.dumps(response_data)}\n\n"
+        except Exception as e:
+            error_data = {"error": str(e), "type": "error"}
+            yield f"data: {json.dumps(error_data)}\n\n"
     
-    return jsonify({
-        "response": response,
-        "time": end-start
-    })
+    return Response(
+        stream_with_context(generate()),
+        mimetype='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive'
+        }
+    )
 
 @app.route("/query", methods=["POST"])
 def query_model():
-    global audio_telugu_response
     data = request.get_json()
-    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    print("Entered '/query' at time: ",current_time)
     query = data.get("query", "")
-
     if not query:
         return jsonify({"error": "Query cannot be empty"}), 400
-
-    start_eng = time.time()
-    # Assuming `generate_response` is a function that processes the query
-    response = generate_response(model, query)
-    end_eng = time.time()
-    def clean_response(response):
-        return response.replace("<EOS>", "").replace("<SOS>", "").strip()
-    response=clean_response(response)
-    start_te = time.time()
-    telugu_response = translate(model_telugu, response, en_vocab, te_vocab, te_inv_vocab)
-    end_te = time.time()
-    audio_telugu_response=telugu_response
-    return jsonify({"telugu":(telugu_response),"english":(response),"eng_time":(end_eng-start_eng),"telugu_time":(end_te-start_te)})
-
     
+    def generate():
+        try:
+            start = time.time()
+            word_count = 0
+            model.eval()
+            
+            with torch.no_grad():
+                # Generate English response
+                src = torch.tensor(tokenizer.encode(query)).unsqueeze(0).to(device)
+                tgt = torch.tensor([[1]]).to(device)  # < SOS >
+                
+                english_words = []
+                for _ in range(200):  # 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
+                        
+                    word = tokenizer.idx2word.get(next_token.item(), "<UNK>")
+                    if word not in ["<PAD>", "<EOS>", "< SOS >"]:
+                        english_words.append(word.strip())
+                        word_count += 1
+                        response_data = {
+                            "word": word.strip(),
+                            "timestamp": time.time() - start,
+                            "word_count": word_count,
+                            "type": "english"
+                        }
+                        yield f"data: {json.dumps(response_data)}\n\n"
+                
+                # Translate to Telugu if available
+                if translation_available:
+                    english_response = " ".join(english_words)
+                    telugu_response = translate(model_telugu, english_response, en_vocab, te_vocab, te_inv_vocab)
+                    
+                    for word in telugu_response.split():
+                        word_count += 1
+                        response_data = {
+                            "word": word.strip(),
+                            "timestamp": time.time() - start,
+                            "word_count": word_count,
+                            "type": "telugu"
+                        }
+                        yield f"data: {json.dumps(response_data)}\n\n"
+        except Exception as e:
+            error_data = {"error": str(e), "type": "error"}
+            yield f"data: {json.dumps(error_data)}\n\n"
+    
+    return Response(
+        stream_with_context(generate()),
+        mimetype='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive'
+        }
+    )
+
 @app.route("/audio", methods=["POST"])
 def get_audio():
     data = request.get_json()
     text = data.get("text")
-    start_te = time.time()
-
+    
     if not text:
         return jsonify({"error": "No Response To convert to speech"}), 400
 
-    # Convert text to Telugu speech using in-memory file
-    speech = gTTS(text=text, lang="te")
-    audio_io = io.BytesIO()
-    speech.write_to_fp(audio_io)
-    audio_io.seek(0)
-    end_te = time.time()
-    print("telugu_time: ",(end_te-start_te))
-
-    return send_file(audio_io, mimetype="audio/mpeg", as_attachment=False)
+    try:
+        start_te = time.time()
+        # Convert text to Telugu speech using in-memory file
+        speech = gTTS(text=text, lang="te")
+        audio_io = io.BytesIO()
+        speech.write_to_fp(audio_io)
+        audio_io.seek(0)
+        end_te = time.time()
+        print("telugu_time: ", (end_te - start_te))
+
+        return send_file(audio_io, mimetype="audio/mpeg", as_attachment=False)
+    except Exception as e:
+        return jsonify({"error": f"Audio generation failed: {str(e)}"}), 500
+
+if __name__ == "__main__":
+    print("Starting Flask application...")
+    print(f"Translation service available: {translation_available}")
+    app.run(host="0.0.0.0", debug=True)