Update app.py

app.py

@@ -1,94 +1,82 @@
 import gradio as gr
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from torch.optim import AdamW
-from torch.utils.data import Dataset, DataLoader
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torch.optim import Adam
 from opacus import PrivacyEngine
+import os
 
-# ---------------------
-# تنظیمات مدل
-# ---------------------
+# ---------- تنظیمات مدل ----------
 MODEL_NAME = "HooshvareLab/gpt2-fa"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-     torch_dtype="auto"
-)
+MAX_NEW_TOKENS = 150
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
 
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if device=="cuda" else -1)
 
-# ---------------------
-# Dataset ساده
-# ---------------------
-class TextDataset(Dataset):
-    def __init__(self, file_path, tokenizer, max_length=128):
-        with open(file_path, "r", encoding="utf-8") as f:
-            lines = f.read().splitlines()
-        self.examples = [tokenizer.encode(line, truncation=True, max_length=max_length) for line in lines]
+# ---------- داده آموزشی ----------
+TRAINING_FILE = "training_data.txt"
 
-    def __len__(self):
-        return len(self.examples)
+if not os.path.exists(TRAINING_FILE):
+    # ایجاد یک فایل نمونه کوچک
+    with open(TRAINING_FILE, "w", encoding="utf-8") as f:
+        f.write("سلام! این یک متن نمونه برای آموزش مدل GPT-2 فارسی است.\n")
+        f.write("مدل می‌تواند پاسخ‌های ساده و امن بدهد.\n")
+
+with open(TRAINING_FILE, "r", encoding="utf-8") as f:
+    texts = f.read().splitlines()
 
+class TextDataset(Dataset):
+    def __init__(self, texts, tokenizer, max_length=128):
+        self.encodings = tokenizer(texts, truncation=True, padding=True, max_length=max_length, return_tensors="pt")
+    def __len__(self):
+        return len(self.encodings["input_ids"])
     def __getitem__(self, idx):
-        return torch.tensor(self.examples[idx])
+        return {key: val[idx] for key, val in self.encodings.items()}
 
-# بارگذاری داده آموزشی
-train_dataset = TextDataset("training_data.txt", tokenizer)
-train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
+dataset = TextDataset(texts, tokenizer)
+dataloader = DataLoader(dataset, batch_size=2, shuffle=True)
 
-# ---------------------
-# تنظیمات بهینه‌ساز و DP
-# ---------------------
-optimizer = AdamW(model.parameters(), lr=5e-5)
-privacy_engine = PrivacyEngine(
+# ---------- بهینه‌ساز و DP ----------
+optimizer = Adam(model.parameters(), lr=5e-5)
+
+privacy_engine = PrivacyEngine()
+privacy_engine.attach(
     model,
-    batch_size=2,
-    sample_size=len(train_dataset),
-    alphas=[10, 100],
-    noise_multiplier=1.0,
-    max_grad_norm=1.0,
+    optimizer,
+    sample_rate=dataloader.batch_size / len(dataset),
+    max_grad_norm=1.0
 )
-privacy_engine.attach(optimizer)
-
-# ---------------------
-# تابع تولید متن
-# ---------------------
-def generate_and_learn(prompt, max_new_tokens=150, temperature=0.7, top_p=0.9, train_steps=1):
-    # تولید متن
-    output_text = generator(prompt, max_new_tokens=int(max_new_tokens),
-                            do_sample=True, temperature=float(temperature),
-                            top_p=float(top_p))[0]["generated_text"]
 
-    # یادگیری incremental روی داده آموزشی
+# ---------- تابع آموزش incremental ----------
+def train_one_epoch():
     model.train()
-    for _ in range(train_steps):
-        for batch in train_loader:
-            batch = batch.to(device)
-            outputs = model(batch, labels=batch)
-            loss = outputs.loss
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+    for batch in dataloader:
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+        outputs = model(input_ids, attention_mask=attention_mask, labels=input_ids)
+        loss = outputs.loss
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+    return f"✅ یک دوره آموزشی انجام شد. آخرین Loss: {loss.item():.4f}"
 
-    return output_text + "\n\n[✅ مدل بعد از یادگیری incremental بروزرسانی شد]"
+# ---------- تابع تولید متن ----------
+def generate_text(prompt, max_new_tokens=MAX_NEW_TOKENS):
+    model.eval()
+    try:
+        return generator(prompt, max_new_tokens=int(max_new_tokens), do_sample=True)[0]["generated_text"]
+    except torch.cuda.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        return "⚠️ حافظه کم بود، خروجی کامل تولید نشد."
+    except Exception as e:
+        return f"❌ خطا: {e}"
 
-# ---------------------
-# رابط Gradio
-# ---------------------
+# ---------- رابط Gradio ----------
 with gr.Blocks(title="چت بات فارسی با یادگیری امن") as demo:
-    gr.Markdown(
-        "<div dir='rtl'>"
-        "## چت بات فارسی با یادگیری امن\n"
-        "یک Prompt وارد کنید و پاسخ مدل را دریافت کنید. مدل قابلیت یادگیری incremental با **DP** دارد."
-        "</div>"
-    )
-    prompt_box = gr.Textbox(label="Prompt", lines=4)
-    btn = gr.Button("تولید پاسخ")
-    output_box = gr.Textbox(label="خروجی", lines=12)
-
-    btn.click(generate_and_learn, inputs=prompt_box, outputs=output_box)
-
-if __name__ == "__main__":
-    demo.launch()
+    gr.Markdown("<div dir='rtl'>## چت بات فارسی (DP-ready)\n"
+                "تنها یک پرامپت وارد کنید و پاسخ مدل را دریافت کنید. "
+                "مدل به صورت incremental روی داده‌های داخلی آموزش داده می‌شود.</div>", elem_id="int_