Update app.py

app.py

@@ -1,35 +1,58 @@
 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)
 
 url = "https://drive.google.com/uc?id=1RCZShB5ohy1HdU-mogcP16TbeVv9txpY"
 df = pd.read_csv(url)
-# Tokenizer
+
+# Optimized Tokenizer with caching
 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
+        self._encode_cache = {}  # Cache for faster encoding
 
     def build_vocab(self, texts):
+        # Optimized vocabulary building
+        unique_words = set()
         for text in texts:
-            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
+            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
 
     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]
-        return tokens + [0] * (max_len - len(tokens))
+        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
 
     def decode(self, tokens):
         return " ".join([self.idx2word.get(idx, "<UNK>") for idx in tokens if idx > 0])
@@ -41,86 +64,246 @@ 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())
 
-# Model
+# Optimized Model for HF Spaces
 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))
-        self.transformer = nn.TransformerDecoder(
-            nn.TransformerDecoderLayer(d_model=embed_size, nhead=num_heads),
-            num_layers=num_layers
+        
+        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(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("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")  # HF Spaces typically CPU-only
 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):
-        model.load_state_dict(torch.load(path, map_location=device))
+        # 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.eval()
         print("Model loaded successfully.")
     else:
-        print("Model file not found!")
-
+        print("Model file not found! Using randomly initialized model.")
 
-def generate_response_stream(model, query, max_length=200):
+# 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"""
     model.eval()
+    
     with torch.no_grad():
-        src = torch.tensor(tokenizer.encode(query)).unsqueeze(0).to(device)
-        tgt = torch.tensor([[1]]).to(device)  # <SOS>
+        # 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
         
-        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)
-            
-            # Get the current word
-            current_word = tokenizer.idx2word.get(next_token.item(), "<UNK>")
-            if current_word != "<PAD>":
-                yield current_word + " "
-            
-            if next_token.item() == 2:  # <EOS>
+        words_buffer = []
+        consecutive_repeats = 0
+        last_word = ""
+        
+        for step in range(max_length):
+            try:
+                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)
+                
+                tgt = torch.cat([tgt, next_token], dim=1)
+                
+                token_id = next_token.item()
+                if token_id == 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
+                
+                words_buffer.append(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()
 
-# Flask App
+# Flask App optimized for HF Spaces
 app = Flask(__name__)
 
 @app.route("/")
 def home():
-    return {"message": "Streaming Transformer-based Response Generator API is running!"}
+    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}
 
 @app.route("/intent")
 def intents():
-    return jsonify({"intents": list(set(df['intent'].dropna()))})
+    try:
+        return jsonify({"intents": list(set(df['intent'].dropna()))})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
 @app.route("/query", methods=["POST"])
 def query_model():
-    data = request.get_json()
-    query = data.get("query", "")
-    if not query:
-        return jsonify({"error": "Query cannot be empty"}), 400
-    
-    def generate():
-        start = time.time()
-        for word in generate_response_stream(model, query):
-            response_data = {
-                "word": word,
-                "timestamp": time.time() - start
+    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': '*'
             }
-            yield f"data: {json.dumps(response_data)}\n\n"
-    
-    return Response(stream_with_context(generate()), mimetype='text/event-stream')
+        )
+        
+    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
 
 if __name__ == "__main__":
+    print("Loading model...")
     load_model(model)
-    app.run(host="0.0.0.0", port=7860) 
+    print("Starting HF Spaces optimized server...")
+    
+    # HF Spaces compatible settings
+    port = int(os.environ.get("PORT", 7860))
+    app.run(
+        host="0.0.0.0", 
+        port=port, 
+        debug=False,  # Disable debug for production
+        threaded=True
+    )