Update app.py

app.py

@@ -1,58 +1,43 @@
 import torch
 import torch.nn as nn
+import torch.optim as optim
 import pandas as pd
+from torch.utils.data import Dataset, DataLoader
 from flask import Flask, request, jsonify, Response, stream_with_context
 from sklearn.model_selection import train_test_split
 import os
 import time
 import json
-import threading
-from queue import Queue
-import multiprocessing
 
-# Optimize for Hugging Face Spaces CPU limits
-num_cores = min(multiprocessing.cpu_count(), 4)  # HF Spaces usually have 2-4 cores
-torch.set_num_threads(num_cores)
-torch.set_num_interop_threads(num_cores)
+# Set PyTorch to use all available CPU threads
+torch.set_num_threads(os.cpu_count())
+torch.set_num_interop_threads(os.cpu_count())
+
+# Enable optimizations
+torch.backends.mkldnn.enabled = True if hasattr(torch.backends, 'mkldnn') else False
 
 url = "https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY"
 df = pd.read_csv(url)
 
-# Optimized Tokenizer with caching
+# Tokenizer
 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.vocab_size = 4
-        self._encode_cache = {}  # Cache for faster encoding
 
     def build_vocab(self, texts):
-        # Optimized vocabulary building
-        unique_words = set()
         for text in texts:
-            unique_words.update(text.split())
-        
-        for word in sorted(unique_words):  # Sort for consistent ordering
-            if word not in self.word2idx:
-                self.word2idx[word] = self.vocab_size
-                self.idx2word[self.vocab_size] = word
-                self.vocab_size += 1
+            for word in text.split():
+                if word not in self.word2idx:
+                    self.word2idx[word] = self.vocab_size
+                    self.idx2word[self.vocab_size] = word
+                    self.vocab_size += 1
 
     def encode(self, text, max_len=200):
-        # Use cache for repeated queries
-        cache_key = (text, max_len)
-        if cache_key in self._encode_cache:
-            return self._encode_cache[cache_key]
-            
         tokens = [self.word2idx.get(word, 3) for word in text.split()]
         tokens = [1] + tokens[:max_len - 2] + [2]
-        encoded = tokens + [0] * (max_len - len(tokens))
-        
-        # Cache result
-        if len(self._encode_cache) < 1000:  # Limit cache size
-            self._encode_cache[cache_key] = encoded
-            
-        return encoded
+        return tokens + [0] * (max_len - len(tokens))
 
     def decode(self, tokens):
         return " ".join([self.idx2word.get(idx, "<UNK>") for idx in tokens if idx > 0])
@@ -64,246 +49,131 @@ train_data, test_data = train_test_split(df, test_size=0.2, random_state=42)
 tokenizer = ScratchTokenizer()
 tokenizer.build_vocab(train_data["instruction"].tolist() + train_data["response"].tolist())
 
-# Optimized Model for HF Spaces
+# Model
 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__()
-        # Reduced model size for HF Spaces memory limits
         self.embedding = nn.Embedding(vocab_size, embed_size)
         self.pos_embedding = nn.Parameter(torch.randn(1, max_len, embed_size))
-        
-        decoder_layer = nn.TransformerDecoderLayer(
-            d_model=embed_size, 
-            nhead=num_heads,
-            dim_feedforward=embed_size * 2,  # Reduced from 4x to 2x
-            dropout=0.1,
-            activation='gelu',
-            batch_first=False,
-            norm_first=True  # Pre-norm for better stability
+        self.transformer = nn.TransformerDecoder(
+            nn.TransformerDecoderLayer(d_model=embed_size, nhead=num_heads),
+            num_layers=num_layers
         )
-        
-        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
         self.fc_out = nn.Linear(embed_size, vocab_size)
-        self.max_len = max_len
 
     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), :]
-        
         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
-        )
+        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))
 
 # Load model
-device = torch.device("cpu")  # HF Spaces typically CPU-only
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = GPTModel(tokenizer.vocab_size).to(device)
 
-# Try to optimize with torch.jit if available
-try:
-    # Create a traced model for faster inference
-    sample_src = torch.randint(0, tokenizer.vocab_size, (1, 50))
-    sample_tgt = torch.randint(0, tokenizer.vocab_size, (1, 10))
-    traced_model = torch.jit.trace(model, (sample_src, sample_tgt))
-    model = traced_model
-    print("Model traced with TorchScript for faster inference")
-except Exception as e:
-    print(f"TorchScript tracing failed: {e}, using regular model")
-
 def load_model(model, path="gpt_model.pth"):
     if os.path.exists(path):
-        # Load with CPU mapping for HF Spaces
-        checkpoint = torch.load(path, map_location='cpu')
-        
-        # Handle different checkpoint formats
-        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
-            model.load_state_dict(checkpoint['state_dict'])
-        else:
-            model.load_state_dict(checkpoint)
-            
+        model.load_state_dict(torch.load(path, map_location=device, weights_only=True))
         model.eval()
+        # Enable inference optimizations
+        if hasattr(torch.jit, 'optimize_for_inference'):
+            model = torch.jit.optimize_for_inference(torch.jit.script(model))
         print("Model loaded successfully.")
     else:
-        print("Model file not found! Using randomly initialized model.")
+        print("Model file not found!")
+    return model
 
-# Optimized generation with batching and early stopping
-def generate_response_stream_fast(model, query, max_length=200, chunk_size=3):
-    """Optimized generation for HF Spaces"""
+def generate_response_stream(model, query, max_length=200):
     model.eval()
     
+    # Pre-encode the query once
+    src_tokens = tokenizer.encode(query)
+    src = torch.tensor(src_tokens).unsqueeze(0).to(device)
+    tgt = torch.tensor([[1]], dtype=torch.long).to(device)  # < SOS >
+    
+    # Pre-allocate tensor for better memory efficiency
+    max_tgt_len = min(max_length, 200)
+    
     with torch.no_grad():
-        # Use smaller sequences for HF Spaces
-        src = torch.tensor(tokenizer.encode(query, max_len=200)).unsqueeze(0).to(device)
-        tgt = torch.tensor([[1]]).to(device)  # SOS token
-        
-        words_buffer = []
-        consecutive_repeats = 0
-        last_word = ""
-        
-        for step in range(max_length):
-            try:
+        # Use torch.inference_mode for better performance
+        with torch.inference_mode():
+            for step in range(max_length):
+                # Forward pass
                 output = model(src, tgt)
                 
-                # Use top-k sampling instead of greedy for better responses
-                logits = output[:, -1, :] / 0.8  # Temperature scaling
-                top_k = torch.topk(logits, k=5)
-                probs = torch.softmax(top_k.values, dim=-1)
-                next_token_idx = torch.multinomial(probs, 1)
-                next_token = top_k.indices.gather(-1, next_token_idx)
+                # Get next token more efficiently
+                logits = output[:, -1, :]
+                next_token = torch.argmax(logits, dim=-1, keepdim=True)
                 
-                tgt = torch.cat([tgt, next_token], dim=1)
-                
-                token_id = next_token.item()
-                if token_id == 2:  # EOS
+                # Check for EOS early
+                if next_token.item() == 2:  # <EOS>
                     break
-                    
-                word = tokenizer.idx2word.get(token_id, "<UNK>")
                 
-                # Skip special tokens and repeated words
-                if word in ["<PAD>", "< SOS >", "<EOS>", "<UNK>"]:
-                    continue
-                    
-                # Prevent infinite loops
-                if word == last_word:
-                    consecutive_repeats += 1
-                    if consecutive_repeats > 2:
-                        continue
-                else:
-                    consecutive_repeats = 0
-                    last_word = word
+                # Concatenate token
+                tgt = torch.cat([tgt, next_token], dim=1)
                 
-                words_buffer.append(word)
+                # Get the current word
+                current_word = tokenizer.idx2word.get(next_token.item(), "<UNK>")
+                if current_word not in ["<PAD>", "<EOS>", "< SOS >"]:
+                    yield current_word + " "
                 
-                # Stream in chunks for better perceived performance
-                if len(words_buffer) >= chunk_size:
-                    chunk_text = " ".join(words_buffer) + " "
-                    words_buffer = []
-                    yield chunk_text
-                    
-            except Exception as e:
-                print(f"Generation error at step {step}: {e}")
-                break
-        
-        # Yield remaining words
-        if words_buffer:
-            yield " ".join(words_buffer) + " "
-
-# Simple request queue for better CPU utilization
-request_queue = Queue(maxsize=10)
-processing_lock = threading.Lock()
+                # Prevent infinite loops
+                if tgt.size(1) >= max_tgt_len:
+                    break
 
-# Flask App optimized for HF Spaces
+# Flask App with threading optimizations
 app = Flask(__name__)
 
+# Configure Flask for better performance
+app.config['THREADED'] = True
+
 @app.route("/")
 def home():
-    return {
-        "message": "HF Spaces Optimized Transformer API", 
-        "status": "running",
-        "device": str(device),
-        "vocab_size": tokenizer.vocab_size
-    }
-
-@app.route("/health")
-def health():
-    return {"status": "healthy", "model_loaded": True}
+    return {"message": "Streaming Transformer-based Response Generator API is running!"}
 
 @app.route("/intent")
 def intents():
-    try:
-        return jsonify({"intents": list(set(df['intent'].dropna()))})
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+    return jsonify({"intents": list(set(df['intent'].dropna()))})
 
 @app.route("/query", methods=["POST"])
 def query_model():
-    try:
-        data = request.get_json()
-        query = data.get("query", "").strip()
-        
-        if not query:
-            return jsonify({"error": "Query cannot be empty"}), 400
-            
-        if len(query) > 500:  # Limit input length for HF Spaces
-            query = query[:500]
-        
-        def generate():
-            start_time = time.time()
-            word_count = 0
-            
-            try:
-                for chunk in generate_response_stream_fast(model, query, max_length=50):
-                    word_count += len(chunk.split())
-                    response_data = {
-                        "chunk": chunk,
-                        "timestamp": time.time() - start_time,
-                        "word_count": word_count
-                    }
-                    yield f"data: {json.dumps(response_data)}\n\n"
-                    
-                    # Prevent very long responses on HF Spaces
-                    if word_count > 100:
-                        break
-                        
-            except Exception as e:
-                error_data = {
-                    "error": f"Generation failed: {str(e)}",
-                    "timestamp": time.time() - start_time
-                }
-                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',
-                'Access-Control-Allow-Origin': '*'
+    data = request.get_json()
+    query = data.get("query", "")
+    if not query:
+        return jsonify({"error": "Query cannot be empty"}), 400
+    
+    def generate():
+        start = time.time()
+        word_count = 0
+        for word in generate_response_stream(model, query):
+            word_count += 1
+            response_data = {
+                "word": word.strip(),
+                "timestamp": time.time() - start,
+                "word_count": word_count
             }
-        )
-        
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
-
-@app.route("/simple_query", methods=["POST"])
-def simple_query():
-    """Non-streaming endpoint for simpler clients"""
-    try:
-        data = request.get_json()
-        query = data.get("query", "").strip()
-        
-        if not query:
-            return jsonify({"error": "Query cannot be empty"}), 400
-        
-        start_time = time.time()
-        response_text = ""
-        
-        for chunk in generate_response_stream_fast(model, query, max_length=50):
-            response_text += chunk
-            
-        return jsonify({
-            "query": query,
-            "response": response_text.strip(),
-            "processing_time": time.time() - start_time
-        })
-        
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+            yield f"data: {json.dumps(response_data)}\n\n"
+    
+    return Response(
+        stream_with_context(generate()), 
+        mimetype='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive',
+            'X-Accel-Buffering': 'no'  # Disable nginx buffering if present
+        }
+    )
 
 if __name__ == "__main__":
-    print("Loading model...")
-    load_model(model)
-    print("Starting HF Spaces optimized server...")
+    # Load and optimize model
+    model = load_model(model)
     
-    # HF Spaces compatible settings
-    port = int(os.environ.get("PORT", 7860))
+    # Run Flask with threading enabled and optimized worker settings
     app.run(
         host="0.0.0.0", 
-        port=port, 
-        debug=False,  # Disable debug for production
-        threaded=True
+        port=7860,
+        threaded=True,
+        processes=1,  # Use threading instead of multiprocessing for better memory sharing
+        debug=False   # Disable debug mode for better performance
     )