Update main.py

main.py

@@ -1,31 +1,61 @@
+import os
 import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct")
 
+app = FastAPI(title="Qwen FastAPI")
+
+tokenizer = None
+model = None
+
+
+class GenerateRequest(BaseModel):
+    system_prompt: str
+    user_prompt: str
+    max_new_tokens: int = 400
+    temperature: float = 0.7
+    top_p: float = 0.9
+    do_sample: bool = True
+
+
+@app.on_event("startup")
+def startup_event():
+    global tokenizer, model
 
-def main():
     # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-    # Pick dtype مناسب: bfloat16 لو GPU متاح، غير كده float32 على CPU
+    # dtype: bfloat16 on CUDA, float32 on CPU
     has_cuda = torch.cuda.is_available()
     dtype = torch.bfloat16 if has_cuda else torch.float32
 
-    # Load model (device_map="auto" يوزع تلقائي)
+    # Load model (auto device placement)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         torch_dtype=dtype,
         device_map="auto"
     )
 
-    # Prompt: explain Past Simple in simple English
+    print("Model ready")  # ✅ مطلوب منك
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_NAME}
+
+
+@app.post("/generate")
+def generate(req: GenerateRequest):
+    global tokenizer, model
+
     messages = [
-        {"role": "system", "content": "You are a friendly English teacher. Explain clearly and simply."},
-        {"role": "user", "content": "Explain the Past Simple tense in very simple English. Give rules and 8 short examples. Keep it clear for A2 learners."}
+        {"role": "system", "content": req.system_prompt},
+        {"role": "user", "content": req.user_prompt}
     ]
 
-    # Convert chat messages to model input
     text = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -34,23 +64,24 @@ def main():
 
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    # Generate
+    print("\n=== Incoming Request ===")
+    print("SYSTEM:", req.system_prompt)
+    print("USER:", req.user_prompt)
+
     with torch.no_grad():
         generated_ids = model.generate(
             **model_inputs,
-            max_new_tokens=400,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
+            max_new_tokens=req.max_new_tokens,
+            do_sample=req.do_sample,
+            temperature=req.temperature,
+            top_p=req.top_p,
         )
 
-    # Keep only the newly generated tokens (remove the prompt tokens)
     new_tokens = generated_ids[0, model_inputs["input_ids"].shape[-1]:]
-    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
 
-    print("\n=== Model Response ===\n")
-    print(response.strip())
-    print("\n======================\n")
+    print("\n=== Model Response ===")
+    print(response)
+    print("======================\n")
 
-if __name__ == "__main__":
-    main()
+    return {"response": response}