Update app.py

app.py

@@ -3,66 +3,55 @@ import torch
 from fastapi import FastAPI, Request
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
-from typing import Dict, Any
 import logging
 
 # Log ayarları
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Ortam değişkenleri ve konfigürasyon
-HF_TOKEN = os.getenv("HF_TOKEN")
+# HF Spaces otomatik olarak HF_TOKEN sağlar
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 if not HF_TOKEN:
-    logger.error("HF_TOKEN environment variable not set!")
-    raise ValueError("HF_TOKEN environment variable not set!")
+    logger.warning("HF_TOKEN bulunamadı! Genel modellerle çalışılacak")
 
-# Model konfigürasyonu
+# Model konfigürasyonu (HF Spaces için optimize)
 MODEL_CONFIG = {
     "base_model": "google/gemma-1.1-2b-it",
     "lora_model": "programci48/heytak-lora-v1",
     "cache_dir": "/tmp/huggingface",
+    "offload_folder": "/tmp/offload",
     "device": "cuda" if torch.cuda.is_available() else "cpu",
-    "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
-    "offload_folder": "/tmp/offload"  # Offload için yeni dizin
+    "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32
 }
 
-def load_models() -> Dict[str, Any]:
-    """Modelleri yükleyen fonksiyon"""
+def load_models():
+    """HF Spaces için optimize edilmiş model yükleme"""
     try:
-        # Offload dizini oluştur
-        os.makedirs(MODEL_CONFIG["offload_folder"], exist_ok=True)
-        
-        logger.info("Tokenizer yükleniyor...")
+        # Tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_CONFIG["base_model"],
             token=HF_TOKEN,
             cache_dir=MODEL_CONFIG["cache_dir"]
         )
 
-        logger.info(f"Temel model yükleniyor ({MODEL_CONFIG['device']})...")
+        # Model
         base_model = AutoModelForCausalLM.from_pretrained(
             MODEL_CONFIG["base_model"],
             torch_dtype=MODEL_CONFIG["torch_dtype"],
             device_map="auto" if MODEL_CONFIG["device"] == "cuda" else None,
             token=HF_TOKEN,
-            low_cpu_mem_usage=True,
             cache_dir=MODEL_CONFIG["cache_dir"],
             offload_folder=MODEL_CONFIG["offload_folder"]
         )
 
-        logger.info("LoRA adaptörü yükleniyor...")
+        # LoRA Adaptörü
         model = PeftModel.from_pretrained(
             base_model,
             MODEL_CONFIG["lora_model"],
             token=HF_TOKEN
         )
         model.eval()
-        
-        if MODEL_CONFIG["device"] == "cpu":
-            model = model.to("cpu")
-            torch.cuda.empty_cache()
 
-        logger.info("Modeller başarıyla yüklendi!")
         return {"tokenizer": tokenizer, "model": model}
 
     except Exception as e:
@@ -70,54 +59,53 @@ def load_models() -> Dict[str, Any]:
         raise
 
 # Uygulama başlatma
-try:
-    models = load_models()
-    app = FastAPI(title="Gemma-LoRA API", version="1.0")
-except Exception as e:
-    logger.critical(f"Uygulama başlatılamadı: {str(e)}")
-    raise
-
-# API Endpoint'leri
-@app.post("/run/predict")
+app = FastAPI(title="HeyTak AI API")
+
+@app.on_event("startup")
+async def startup_event():
+    try:
+        app.state.models = load_models()
+        logger.info("Modeller başarıyla yüklendi!")
+    except Exception as e:
+        logger.critical(f"Başlatma hatası: {str(e)}")
+        raise
+
+@app.post("/predict")
 async def predict(request: Request):
     try:
         data = await request.json()
-        prompt = data["data"][0]
-        logger.info(f"Gelen istek: {prompt[:50]}...")
-
-        inputs = models["tokenizer"](
-            prompt,
+        prompt = data.get("inputs", "")
+        
+        inputs = app.state.models["tokenizer"](
+            prompt, 
             return_tensors="pt",
             truncation=True,
             max_length=512
-        ).to(models["model"].device)
+        ).to(app.state.models["model"].device)
 
         with torch.no_grad():
-            outputs = models["model"].generate(
+            outputs = app.state.models["model"].generate(
                 **inputs,
                 max_new_tokens=100,
-                do_sample=True,
                 temperature=0.7,
-                top_p=0.9,
-                repetition_penalty=1.1
+                top_p=0.9
             )
 
-        response = models["tokenizer"].decode(
-            outputs[0],
+        response = app.state.models["tokenizer"].decode(
+            outputs[0], 
             skip_special_tokens=True
         ).strip()
 
-        logger.info(f"Oluşturulan yanıt: {response[:50]}...")
-        return {"data": [response]}
+        return {"generated_text": response}
 
     except Exception as e:
-        logger.error(f"İşlem hatası: {str(e)}")
+        logger.error(f"Tahmin hatası: {str(e)}")
         return {"error": str(e)}, 500
 
-@app.get("/health")
+@app.get("/")
 async def health_check():
     return {
-        "status": "healthy",
+        "status": "active",
         "device": MODEL_CONFIG["device"],
-        "torch_dtype": str(MODEL_CONFIG["torch_dtype"])
+        "framework": "FastAPI"
     }