Update app.py

app.py

@@ -3,40 +3,103 @@ import torch
 from fastapi import FastAPI, Request
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+from huggingface_hub import login
+from typing import Dict, Any
 
-# Hugging Face token (gated modeller için gerekli)
-hf_token = os.getenv("HF_TOKEN")
-print("HF_TOKEN:", hf_token)  # Çıktı logs içinde görünür
+# Hugging Face token
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable not set!")
 
-# Model ID'leri
-base_model_id = "google/gemma-1.1-2b-it"
-lora_model_id = "programci48/heytak-lora-v1"
+# Login to Hugging Face Hub
+login(token=HF_TOKEN)
 
-# Tokenizer ve model yükleme
-tokenizer = AutoTokenizer.from_pretrained(base_model_id, token=hf_token)
+# Model IDs
+BASE_MODEL_ID = "google/gemma-1.1-2b-it"
+LORA_MODEL_ID = "programci48/heytak-lora-v1"
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,
-    torch_dtype=torch.float32,
-    device_map=None,  # Hugging Face CPU ortamı için GPU ayarı yapılmaz
-    token=hf_token
-)
+# Load models with error handling and optimizations
+def load_models() -> Dict[str, Any]:
+    try:
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            BASE_MODEL_ID,
+            token=HF_TOKEN
+        )
 
-model = PeftModel.from_pretrained(base_model, lora_model_id, token=hf_token)
-model.eval()
+        # Load base model with memory optimization
+        base_model = AutoModelForCausalLM.from_pretrained(
+            BASE_MODEL_ID,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto",
+            token=HF_TOKEN,
+            low_cpu_mem_usage=True,
+            offload_folder="offload"  # For CPU offloading if needed
+        )
 
-# FastAPI uygulaması
-app = FastAPI()
+        # Load LoRA adapter
+        model = PeftModel.from_pretrained(
+            base_model,
+            LORA_MODEL_ID,
+            token=HF_TOKEN
+        )
+        model.eval()
+
+        # Move to CPU if no GPU available
+        if not torch.cuda.is_available():
+            model = model.to("cpu")
+            print("Model moved to CPU")
+
+        return {
+            "tokenizer": tokenizer,
+            "model": model
+        }
+
+    except Exception as e:
+        raise RuntimeError(f"Model loading failed: {str(e)}")
+
+# Initialize models
+models = load_models()
+
+# FastAPI app
+app = FastAPI(title="Gemma-LoRA API")
 
 @app.post("/run/predict")
 async def predict(request: Request):
-    data = await request.json()
-    prompt = data["data"][0]
+    try:
+        data = await request.json()
+        prompt = data["data"][0]
+
+        # Tokenize with truncation
+        inputs = models["tokenizer"](
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        ).to(models["model"].device)
+
+        # Generate response
+        with torch.no_grad():
+            outputs = models["model"].generate(
+                **inputs,
+                max_new_tokens=100,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.1
+            )
+
+        # Decode and clean response
+        response = models["tokenizer"].decode(
+            outputs[0],
+            skip_special_tokens=True
+        ).strip()
+
+        return {"data": [response]}
 
-    # Model ile yanıt üret
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_new_tokens=100)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        return {"error": str(e)}, 500
 
-    return {"data": [response]}
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}