Update app.py

app.py

@@ -1,79 +1,88 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from torch.utils.data import DataLoader, Dataset
+from torch.nn import CrossEntropyLoss
+from opacus import PrivacyEngine
 
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  
-
+MODEL_NAME = "HooshvareLab/gpt2-fa"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",           
-    torch_dtype="auto"           
-)
-
-
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-def generate_text(prompt, max_new_tokens, temperature, top_p):
-    """
- میتونی حرف بزنی
-    """
-    max_new_tokens = int(max_new_tokens)
-    temperature = float(temperature)
-    top_p = float(top_p)
-
-    def _run(tokens):
-        return generator(
-            prompt,
-            max_new_tokens=tokens,  # مستقل از طول ورودی
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-        )[0]["generated_text"]
-
-    try:
-        return _run(max_new_tokens)
-    except torch.cuda.OutOfMemoryError:
-        # تلاش مجدد با نصف مقدار
-        try:
-            torch.cuda.empty_cache()
-            smaller = max(32, max_new_tokens // 2)
-            return _run(smaller) + "\n\n[⚠️ حافظه کم بود؛ خروجی با توکن‌های کمتر تولید شد. اگر باز هم مشکل داری، مقدار max_new_tokens را کمتر کن یا از مدل کوچک‌تر استفاده کن.]"
-        except Exception:
-            torch.cuda.empty_cache()
-            return "⚠️ خطای حافظه (OOM). لطفاً مقدار max_new_tokens را کاهش بده یا از مدل کوچک‌تر/سخت‌افزار قوی‌تر استفاده کن."
-    except RuntimeError as e:
-        # برخی OOMها RuntimeError هستند
-        if "out of memory" in str(e).lower():
-            try:
-                torch.cuda.empty_cache()
-                smaller = max(32, max_new_tokens // 2)
-                return _run(smaller) + "\n\n[⚠️ حافظه کم بود؛ خروجی با توکن‌های کمتر تولید شد. اگر باز هم مشکل داری، مقدار max_new_tokens را کمتر کن یا از مدل کوچک‌تر استفاده کن.]"
-            except Exception:
-                torch.cuda.empty_cache()
-                return "⚠️ خطای حافظه (OOM). لطفاً مقدار max_new_tokens را کاهش بده یا از مدل کوچک‌تر/سخت‌افزار قوی‌تر استفاده کن."
-        # سایر خطاها
-        return f"❌ خطا: {e}"
-    except Exception as e:
-        return f"❌ خطا: {e}"
-
-with gr.Blocks(title="TinyLlama Chat") as demo:
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+
+# ---------- Dataset ساده برای PyTorch ----------
+class UserDataset(Dataset):
+    def __init__(self, texts, tokenizer, max_length=128):
+        self.examples = [tokenizer(t, truncation=True, padding="max_length", max_length=max_length, return_tensors="pt") for t in texts]
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, idx):
+        item = self.examples[idx]
+        return {
+            "input_ids": item["input_ids"].squeeze(0),
+            "attention_mask": item["attention_mask"].squeeze(0)
+        }
+
+
+# ---------- تابع آموزش با DP ----------
+def fine_tune_dp(prompt, user_text, epochs=1, lr=5e-5, max_new_tokens=150, batch_size=1):
+    dataset = UserDataset([user_text], tokenizer)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    criterion = CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
+
+    # ---------- فعال کردن Differential Privacy ----------
+    privacy_engine = PrivacyEngine(
+        model,
+        sample_rate=batch_size / len(dataset),
+        noise_multiplier=1.0,
+        max_grad_norm=1.0,
+    )
+    privacy_engine.attach(optimizer)
+
+    model.train()
+    for epoch in range(epochs):
+        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)
+            logits = outputs.logits
+
+            loss = criterion(logits.view(-1, logits.size(-1)), input_ids.view(-1))
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+    # ---------- تولید متن بعد از آموزش ----------
+    model.eval()
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    with torch.no_grad():
+        output = model.generate(input_ids, max_new_tokens=max_new_tokens, do_sample=True)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+
+# ---------- رابط گریدیو ----------
+with gr.Blocks(title="چت بات فارسی با DP") as demo:
     gr.Markdown(
-        "## TinyLlama Chat\n"
-        "پرامپتت را بنویس و پاسخ بگیر.\n\n"
-        "**نکته:** اگر خطای حافظه دیدی، مقدار `max_new_tokens` را کم کن یا مدل کوچک‌تری انتخاب کن."
+        "## چت بات فارسی با یادگیری امن\n"
+        "متن آموزشی خود را بده و پاسخ بگیر.\n"
+        "گرادیان‌ها با Differential Privacy محافظت می‌شوند؛ داده واقعی کاربر لو نمی‌رود."
     )
-    with gr.Row():
-        prompt = gr.Textbox(label="Prompt", lines=4)
-    with gr.Row():
-        max_new_tokens = gr.Slider(16, 512, value=150, step=1, label="max_new_tokens")
-        temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="temperature")
-        top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
-    btn = gr.Button("تولید پاسخ")
+    prompt = gr.Textbox(label="Prompt", lines=4)
+    user_text = gr.Textbox(label="متن آموزشی برای مدل", lines=4)
+    epochs = gr.Slider(1, 3, value=1, step=1, label="تعداد Epochها")
+    max_new_tokens = gr.Slider(16, 512, value=150, step=1, label="max_new_tokens")
+    btn = gr.Button("تولید پاسخ با یادگیری امن")
     output = gr.Textbox(label="خروجی", lines=12)
-
-    btn.click(generate_text, [prompt, max_new_tokens, temperature, top_p], output)
+    
+    btn.click(fine_tune_dp, [prompt, user_text, epochs, max_new_tokens], output)
 
 if __name__ == "__main__":
     demo.launch()