Update app.py

app.py

@@ -1,9 +1,13 @@
 import os
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    BitsAndBytesConfig
+)
 
-# 🔧 نخلي الكاش في فولدر قابل للكتابة
 os.environ["HF_HOME"] = "/data"
 os.environ["TRANSFORMERS_CACHE"] = "/data/cache"
 
@@ -14,7 +18,6 @@ app = FastAPI(
 
 MODEL_REPO = "Reem1/hakayti-merged"
 
-# مثال توضيحي few-shot
 few_shot_example = """<s>[User]: اسمي كريم وأحب الفضاء. 
 اكتب لي قصة طويلة للأطفال (40 جملة على الأقل) مكتوبة بأسلوب قصصي ممتع.
 [Assistant]:
@@ -26,43 +29,44 @@ few_shot_example = """<s>[User]: اسمي كريم وأحب الفضاء.
 وهكذا تعلم كريم أن الحلم يبدأ بخطوة صغيرة، وأن حب المعرفة يمكن أن يفتح له أبوابًا واسعة نحو المستقبل.
 </s>"""
 
-# تحميل الموديل مرة واحدة
+bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+
+print("⏳ Loading model in 8-bit ...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_REPO,
+    quantization_config=bnb_config,
+    device_map="auto",
+    low_cpu_mem_usage=True
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
+
 pipe = pipeline(
     "text-generation",
-    model=MODEL_REPO,
-    tokenizer=MODEL_REPO,
+    model=model,
+    tokenizer=tokenizer,
     max_new_tokens=900,
     temperature=0.9,
-    top_p=0.95
+    top_p=0.95,
+    do_sample=True
 )
+print("✅ Model loaded.")
 
-# body request
 class StoryRequest(BaseModel):
     name: str
     hobby: str
 
-# مولد القصة
-def generate_story(name, hobby):
+def generate_story(name: str, hobby: str) -> str:
     prompt = f"""{few_shot_example}
 [User]: اسمي {name} وأحب {hobby}. 
 اكتب لي قصة طويلة للأطفال (40 جملة على الأقل) مكتوبة بنفس الأسلوب السابق،
 وتأكد أن القصة كلها تدور حول {hobby} فقط.
 [Assistant]:"""
 
-    out = pipe(
-        prompt,
-        max_new_tokens=900,
-        temperature=0.9,
-        top_p=0.95,
-        do_sample=True
-    )
+    out = pipe(prompt)
     full_text = out[0]["generated_text"]
     story = full_text.split("[Assistant]:")[-1].replace("</s>", "").strip()
     return story
 
-# API endpoint
 @app.post("/generate")
-def generate_story_api(req: StoryRequest):
-    story = generate_story(req.name, req.hobby)
-    return {"story": story}
-
+async def generate_story_api(req: StoryRequest):
+    return {"story": generate_story(req.name, req.hobby)}