Update main.py

main.py

@@ -1,45 +1,74 @@
 import os
+from typing import List, Literal, Optional
+
 import torch
 from fastapi import FastAPI
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct")
+# ----------------------------
+# Model config (matches demo)
+# ----------------------------
+MODEL_NAME = os.getenv("MODEL_NAME", "MBZUAI-Paris/Nile-Chat-12B")
+
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "2024"))
 
-app = FastAPI(title="Qwen FastAPI")
+app = FastAPI(title="Nile-Chat-12B FastAPI")
 
 tokenizer = None
 model = None
 
 
+# ----------------------------
+# Request schemas
+# ----------------------------
+Role = Literal["system", "user", "assistant"]
+
+class ChatMessage(BaseModel):
+    role: Role
+    content: str
+
 class GenerateRequest(BaseModel):
-    system_prompt: str
-    user_prompt: str
-    max_new_tokens: int = 400
-    temperature: float = 0.7
-    top_p: float = 0.9
+    # نفس مفهوم Gradio: history + message
+    # لكن هنا هنوحّدها: messages كاملة، وآخر user message هي الطلب الحالي
+    messages: List[ChatMessage] = Field(..., description="Conversation messages in OpenAI-like format")
+
+    max_new_tokens: int = Field(DEFAULT_MAX_NEW_TOKENS, ge=1, le=MAX_MAX_NEW_TOKENS)
     do_sample: bool = True
+    temperature: float = Field(0.6, ge=0.0, le=4.0)
+    top_p: float = Field(0.9, ge=0.05, le=1.0)
+    top_k: int = Field(50, ge=1, le=1000)
+    repetition_penalty: float = Field(1.1, ge=1.0, le=2.0)
+
 
+class GenerateResponse(BaseModel):
+    response: str
+    trimmed: bool = False
+    model: str = MODEL_NAME
 
+
+# ----------------------------
+# Startup
+# ----------------------------
 @app.on_event("startup")
 def startup_event():
     global tokenizer, model
 
-    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-    # dtype: bfloat16 on CUDA, float32 on CPU
-    has_cuda = torch.cuda.is_available()
-    dtype = torch.bfloat16 if has_cuda else torch.float32
+    # نفس منطق الديمو: bfloat16 + device_map auto
+    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-    # Load model (auto device placement)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
+        device_map="auto",
         torch_dtype=dtype,
-        device_map="auto"
     )
+    model.eval()
 
-    print("Model ready")  # ✅ مطلوب منك
+    print("Model ready")  # ✅ زي ما طلبت
 
 
 @app.get("/health")
@@ -47,41 +76,60 @@ def health():
     return {"status": "ok", "model": MODEL_NAME}
 
 
-@app.post("/generate")
+# ----------------------------
+# Core generation
+# ----------------------------
+@app.post("/generate", response_model=GenerateResponse)
 def generate(req: GenerateRequest):
     global tokenizer, model
 
-    messages = [
-        {"role": "system", "content": req.system_prompt},
-        {"role": "user", "content": req.user_prompt}
-    ]
+    if not req.messages:
+        return GenerateResponse(response="Error: messages is empty", trimmed=False)
+
+    # Nile-Chat demo بيستخدم apply_chat_template على conversation كلها
+    conversation = [m.model_dump() for m in req.messages]
 
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
+    # Build input_ids exactly like the Gradio demo
+    input_ids = tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        return_tensors="pt"
     )
 
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    trimmed = False
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        trimmed = True
+
+    input_ids = input_ids.to(model.device)
 
+    # Logging
+    last_user = next((m.content for m in reversed(req.messages) if m.role == "user"), "")
     print("\n=== Incoming Request ===")
-    print("SYSTEM:", req.system_prompt)
-    print("USER:", req.user_prompt)
+    print("MODEL:", MODEL_NAME)
+    print("LAST USER:", last_user)
+    print("trimmed_input:", trimmed)
+    print("input_tokens:", int(input_ids.shape[1]))
 
+    # Generate (non-streaming API response)
     with torch.no_grad():
-        generated_ids = model.generate(
-            **model_inputs,
+        out = model.generate(
+            input_ids=input_ids,
             max_new_tokens=req.max_new_tokens,
             do_sample=req.do_sample,
-            temperature=req.temperature,
             top_p=req.top_p,
+            top_k=req.top_k,
+            temperature=req.temperature,
+            num_beams=1,
+            repetition_penalty=req.repetition_penalty,
         )
 
-    new_tokens = generated_ids[0, model_inputs["input_ids"].shape[-1]:]
-    response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+    # Decode only new tokens (same idea as your Qwen API)
+    new_tokens = out[0, input_ids.shape[-1]:]
+    response_text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
     print("\n=== Model Response ===")
-    print(response)
+    print(response_text)
     print("======================\n")
 
-    return {"response": response}
+    return GenerateResponse(response=response_text, trimmed=trimmed, model=MODEL_NAME)