Update app.py

app.py

@@ -1,77 +1,137 @@
+import os
+import json
+import asyncio
+import numpy as np
+import onnxruntime as ort
+import tiktoken
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
-import onnxruntime as ort
-import numpy as np
-import tiktoken
-import json
-import os
 
+# App-Initialisierung
 app = FastAPI()
 
-# WICHTIG: Erlaubt deinem externen Frontend den Zugriff
+# CORS für dein externes Frontend
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], # Hier kannst du später deine Domain eintragen
+    allow_origins=["*"],
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Modell & Tokenizer laden
-tokenizer = tiktoken.get_encoding("gpt2")
+# 1. Modell & Tokenizer Setup
+TOKENIZER = tiktoken.get_encoding("gpt2")
 MODEL_PATH = "SmaLLMPro_350M_int8.onnx"
+VOCAB_SIZE = 50304
 
-# Optimierte Session-Optionen für CPU
-sess_options = ort.SessionOptions()
-sess_options.intra_op_num_threads = 2 # HF Spaces haben meist 2 Kerne
-session = ort.InferenceSession(MODEL_PATH, sess_options, providers=['CPUExecutionProvider'])
+# 2. ONNX Runtime Optimierung
+# HF Free Spaces haben 2 vCPUs. Wir limitieren die Threads, 
+# um "Context Switching" Overhead zu vermeiden.
+options = ort.SessionOptions()
+options.intra_op_num_threads = 2
+options.inter_op_num_threads = 2
+options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
 
-def top_k_sample(logits, k=50, temp=0.7):
+# Session laden
+print(f"🚀 Lade Modell {MODEL_PATH} mit CPU-Optimierung...")
+session = ort.InferenceSession(
+    MODEL_PATH, 
+    sess_options=options, 
+    providers=['CPUExecutionProvider']
+)
+
+def fast_top_k_sample(logits, k=25, temp=0.7, penalty=1.2, history=None):
+    """Hochoptimiertes Sampling mit NumPy"""
+    # 1. Repetition Penalty (optional, falls history vorhanden)
+    if history is not None and penalty != 1.0:
+        # Wir bestrafen bereits generierte Tokens direkt in den Logits
+        unique_history = np.unique(history)
+        # Nur gültige Token-Indices bestrafen
+        valid_indices = unique_history[unique_history < len(logits)]
+        logits[valid_indices] /= penalty
+
+    # 2. Temperature Scaling
     logits = logits / max(temp, 1e-6)
-    # Nur die Top-K Werte betrachten (spart massiv Zeit beim Sortieren)
-    top_k_indices = np.argpartition(logits, -k)[-k:]
-    top_k_logits = logits[top_k_indices]
+
+    # 3. Top-K via Partition (schneller als vollständiges Sortieren)
+    # Sucht die k größten Werte ohne den Rest zu sortieren
+    top_k_idx = np.argpartition(logits, -k)[-k:]
+    top_k_logits = logits[top_k_idx]
     
-    # Stabiler Softmax
-    exp_logits = np.exp(top_k_logits - np.max(top_k_logits))
+    # 4. Softmax
+    shifted_logits = top_k_logits - np.max(top_k_logits)
+    exp_logits = np.exp(shifted_logits)
     probs = exp_logits / np.sum(exp_logits)
     
-    return int(np.random.choice(top_k_indices, p=probs))
+    # 5. Sample
+    choice = np.random.choice(top_k_idx, p=probs)
+    return int(choice)
 
 @app.post("/chat")
 async def chat(request: Request):
-    data = await request.json()
-    prompt = f"Instruction:\n{data['prompt']}\n\nResponse:\n"
-    tokens = tokenizer.encode(prompt)
-    
-    max_len = int(data.get('maxLen', 100))
-    temp = float(data.get('temp', 0.7))
-    top_k = int(data.get('topK', 40))
-
-    async def generate():
-        nonlocal tokens
-        for _ in range(max_len):
-            # Kontext auf 1024 beschränken
-            ctx = tokens[-1024:]
-            # Padding (Rechtsbündig)
-            padded = np.zeros((1, 1024), dtype=np.int64)
-            padded[0, -len(ctx):] = ctx
-            
-            # Inferenz
-            outputs = session.run(None, {'input': padded})
-            # Wir nehmen nur die Logits des letzten Tokens
-            logits = outputs[0][0, -1, :50304]
-            
-            next_token = top_k_sample(logits, k=top_k, temp=temp)
-            
-            if next_token == 50256: # EOS
-                break
-            
-            tokens.append(next_token)
-            yield f"data: {json.dumps({'token': tokenizer.decode([next_token])})}\n\n"
-
-    return StreamingResponse(generate(), media_type="text/event-stream")
+    try:
+        data = await request.json()
+        user_prompt = data.get('prompt', '')
+        max_len = int(data.get('maxLen', 100))
+        temp = float(data.get('temp', 0.7))
+        top_k = int(data.get('topK', 25))
+        repetition_penalty = float(data.get('penalty', 1.2))
+
+        # Alpaca Instruction Format
+        full_prompt = f"Instruction:\n{user_prompt}\n\nResponse:\n"
+        tokens = TOKENIZER.encode(full_prompt)
+
+        async def generate():
+            nonlocal tokens
+            # Wir behalten die Historie für die Penalty im Auge
+            history = np.array(tokens, dtype=np.int32)
+
+            for _ in range(max_len):
+                # 1. Context Handling: Immer exakt 1024 (Padding rechtsbündig)
+                ctx = tokens[-1024:]
+                input_array = np.zeros((1, 1024), dtype=np.int64)
+                input_array[0, -len(ctx):] = ctx
+                
+                # 2. Inferenz (Synchroner Call in asynchronem Generator)
+                # Das ist der Flaschenhals, hier arbeitet die CPU
+                outputs = session.run(None, {'input': input_array})
+                
+                # 3. Logits extrahieren (letztes Token, erste VOCAB_SIZE)
+                logits = outputs[0][0, -1, :VOCAB_SIZE].astype(np.float32)
+                
+                # 4. Sampling
+                next_token = fast_top_k_sample(
+                    logits, 
+                    k=top_k, 
+                    temp=temp, 
+                    penalty=repetition_penalty,
+                    history=history
+                )
+                
+                if next_token == 50256: # EOS Token
+                    break
+                
+                # 5. Update state
+                tokens.append(next_token)
+                history = np.append(history, next_token)
+                
+                # 6. Stream zum Client
+                yield f"data: {json.dumps({'token': TOKENIZER.decode([next_token])})}\n\n"
+                
+                # Kurze Pause für den Event-Loop
+                await asyncio.sleep(0.01)
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return {"error": str(e)}
 
 @app.get("/")
-def health():
-    return {"status": "SmaLLMPro API is online"}
+async def health():
+    return {"status": "SmaLLMPro INT8 Engine Online", "threads": options.intra_op_num_threads}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)