Update app.py

app.py

@@ -1,136 +1,75 @@
-# Website: https://madgamesfmai.gamer.gd
-import os
 import torch
-import json
-import uuid
-import base64
-import io
-from PIL import Image
-from fastapi import FastAPI, Request, HTTPException, Depends
-from fastapi.responses import HTMLResponse, JSONResponse
-from fastapi.staticfiles import StaticFiles
-from fastapi.templating import Jinja2Templates
-from transformers import BlipProcessor, BlipForConditionalGeneration
 from model import MiniTransformer, BPETokenizer
-from pydantic import BaseModel
-
+from fastapi import FastAPI, Request, Header, HTTPException
+from huggingface_hub import hf_hub_download
+import uvicorn
+import os
 app = FastAPI()
-templates = Jinja2Templates(directory="templates")
-
-print("Lade Vision-Modell (Augen)...")
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-vision_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cpu")
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 REPO_ID = "Kleinpuki2/madgamesai"
 FILENAME = "madgames_gpt2_stable.pth"
+API_KEY = "MG-ADMIN-1337"
+model = None
 tokenizer = BPETokenizer()
-
-from huggingface_hub import hf_hub_download
-
-print("Lade Modell vom Hugging Face Hub...")
-try:
-    MODEL_PATH = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
-    model = MiniTransformer.load(MODEL_PATH, device=DEVICE)
-    print("Modell erfolgreich geladen!")
-except Exception as e:
-    print(f"Fehler beim Laden des Modells: {e}")
-    model = None
-
-KEYS_FILE = "api_keys.json"
-
-def load_keys():
-    default_keys = {"MG-ADMIN-1337": {"created": "permanent", "uses": 0}}
-    if not os.path.exists(KEYS_FILE): return default_keys
-    with open(KEYS_FILE, "r") as f: 
-        try:
-            keys = json.load(f)
-            keys.update(default_keys)
-            return keys
-        except: return default_keys
-
-def save_keys(keys):
-    with open(KEYS_FILE, "w") as f: json.dump(keys, f)
-
-def verify_api_key(request: Request):
-    api_key = request.headers.get("X-API-Key")
-    keys = load_keys()
-    if api_key not in keys:
-        raise HTTPException(status_code=403, detail="Ungültiger API-Key")
-    return api_key
-
-from pydantic import BaseModel
-from typing import Optional
-
-class PredictRequest(BaseModel):
-    prompt: str
-    image: Optional[str] = None
-
-@app.get("/", response_class=HTMLResponse)
-async def dashboard(request: Request):
-    return templates.TemplateResponse(request=request, name="index.html")
-
-@app.post("/generate_key")
-async def create_key():
-    new_key = f"MG-{uuid.uuid4().hex[:12].upper()}"
-    keys = load_keys()
-    keys[new_key] = {"created": "manual", "uses": 0}
-    save_keys(keys)
-    return {"key": new_key}
-
+def load_model():
+    global model
+    try:
+        print(f"Lade Checkpoint von {REPO_ID}...")
+        path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+        model = MiniTransformer.load(path, device='cpu')
+        print("Modell erfolgreich geladen und initialisiert!")
+        return True
+    except Exception as e:
+        print(f"Fehler beim Laden: {e}")
+        return False
+is_loaded = load_model()
+@app.get("/")
+def root():
+    return {"status": "online", "loaded": is_loaded, "repo": REPO_ID}
 @app.post("/predict")
-async def predict(req: PredictRequest, api_key: str = Depends(verify_api_key)):
+async def predict(request: Request, x_api_key: str = Header(None)):
+    if x_api_key != API_KEY:
+        raise HTTPException(status_code=401)
+    
+    if not is_loaded or model is None:
+        return {"response": "Fehler: Modell konnte nicht korrekt initialisiert werden."}
+    
     try:
-        prompt = req.prompt
+        data = await request.json()
+        prompt = data.get("prompt", "")
+        if not prompt: return {"response": ""}
         
-        if req.image:
-            try:
-                print(f"Verarbeite Bild für Key: {api_key}")
-                img_data = base64.b64decode(req.image.split(",")[-1])
-                img = Image.open(io.BytesIO(img_data)).convert("RGB")
-                inputs = processor(img, return_tensors="pt").to("cpu")
-                out_vision = vision_model.generate(**inputs)
-                description = processor.decode(out_vision[0], skip_special_tokens=True)
-                prompt = f"[Nutzer sendet Bild: {description}] " + prompt
-                print(f"Vision-Ergebnis: {description}")
-            except Exception as e:
-                print(f"Vision-Fehler: {e}")
-
-        if not model:
-            return {"response": "Fehler: Modell wurde nicht korrekt geladen."}
-
-        tokens = tokenizer.encode(prompt)
-        input_ids = torch.tensor([tokens]).to(DEVICE)
-
+        print(f"Anfrage empfangen: {prompt[:50]}...")
+        
+        # WICHTIG: Das exakte Format aus dem Training nachbauen!
+        formatted_prompt = f"User: {prompt}\nKI: "
+        tokens = tokenizer.encode(formatted_prompt)
         ctx_len = model.ctx_len if hasattr(model, 'ctx_len') else 1024
-        input_ids = input_ids[:, -ctx_len:]
+        tokens = tokens[-ctx_len:]
         
-        with torch.no_grad():
-            out = model.generate(input_ids, max_new_tokens=50, temperature=0.2, top_k=10, repetition_penalty=1.0)
+        idx = torch.tensor([tokens], dtype=torch.long)
         
-        generated_tokens = out[0, input_ids.shape[1]:].tolist()
-        response_text = tokenizer.decode(generated_tokens)
+        # Perfekte Settings für Code & Chat
+        out = model.generate(idx, max_new_tokens=250, temperature=0.2, top_k=10, repetition_penalty=1.0)
         
-        response_text = response_text.replace("<|endoftext|>", "").strip()
-        if "User:" in response_text:
-            response_text = response_text.split("User:")[0]
-
-        if not response_text.strip():
-            response_text = "Die KI braucht noch mehr Training. Geh in der Zwischenzeit Credits farmen!"
-        try:
-            keys = load_keys()
-            if api_key in keys:
-                keys[api_key]["uses"] += 1
-                save_keys(keys)
-        except: pass
-            
-        return {"response": response_text}
+        # Nur den neu generierten Teil extrahieren
+        generated_tokens = out[0, len(tokens):].tolist()
+        response = tokenizer.decode(generated_tokens)
+        
+        # <|endoftext|> entfernen, falls es im Text auftaucht
+        response = response.replace("<|endoftext|>", "").strip()
         
+        # Prüfen, ob die KI anfängt, den nächsten "User:" Text zu generieren
+        if "User:" in response:
+            response = response.split("User:")[0]
+            
+        final_text = response.strip()
+        if not final_text:
+            final_text = "Die KI hat noch keine klare Antwort gefunden. Trainiere sie noch ein wenig weiter (Ziel: Loss unter 1.0) oder versuche einen anderen Prompt!"
+            
+        print(f"Antwort generiert: {final_text[:50]}...")
+        return {"response": final_text}
     except Exception as e:
-        print(f"KRITISCHER FEHLER im Predict: {e}")
-        return JSONResponse(status_code=500, content={"error": str(e)})
-
+        print(f"Fehler bei Vorhersage: {e}")
+        return {"response": f"Runtime Fehler: {str(e)}"}
 if __name__ == "__main__":
-    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)
-