Upload app.py

app.py

@@ -1,42 +1,63 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
-import gradio as gr
 import os
 import zipfile
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from peft import PeftModel
+from huggingface_hub import login
 
+# --- گام ۱: احراز هویت Hugging Face
+hf_token = os.environ.get("HF_TOKEN")
+if not hf_token:
+    raise ValueError("❌ HF_TOKEN not found in environment secrets.")
+login(hf_token)
+
+# اگر فایل اشتباه وجود داره و فایل جدید نه
 if os.path.exists("dorna-diabetes-finetuned-20250514T183411Z-1-001.zip") and not os.path.exists("dorna-diabetes-finetuned.zip"):
     os.rename("dorna-diabetes-finetuned-20250514T183411Z-1-001.zip", "dorna-diabetes-finetuned.zip")
     print("✅ اسم فایل تغییر کرد.")
 
+# --- گام ۲: اکسترکت فایل فشرده (فقط بار اول)
 if not os.path.exists("dorna-diabetes-finetuned"):
     with zipfile.ZipFile("dorna-diabetes-finetuned.zip", "r") as zip_ref:
         zip_ref.extractall(".")
-    print("✅ فایل ZIP اکسترکت شد.")
+    print("✅ فایل LoRA اکسترکت شد.")
 
+# --- گام ۳: بارگذاری مدل پایه و LoRA
 BASE_MODEL = "PartAI/Dorna-Llama3-8B-Instruct"
-LORA_PATH = "./dorna-diabetes-finetuned"  # این پوشه رو آپلود می‌کنی توی اسپیس
-
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+LORA_PATH = "./dorna-diabetes-finetuned"
 
+print("🔹 در حال بارگذاری مدل پایه...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_auth_token=hf_token)
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    device_map="auto",
+    load_in_4bit=True,
     torch_dtype=torch.float16,
+    device_map="auto",
+    trust_remote_code=True,
+    use_auth_token=hf_token
 )
 
+print("🔹 در حال بارگذاری LoRA...")
 model = PeftModel.from_pretrained(base_model, LORA_PATH)
+model.eval()
+
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-def generate_response(prompt):
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+# --- گام ۴: رابط چت با Gradio
+def chat(prompt):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
     with torch.no_grad():
         output = model.generate(
             input_ids=input_ids,
             max_new_tokens=200,
-            do_sample=True,
             temperature=0.7,
             top_p=0.9,
+            do_sample=True
         )
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response[len(prompt):].strip()
+
+iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="💬 Dorna LoRA Model")
 
-gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Dorna-Llama3 LoRA").launch()
+iface.launch()