Update app.py

app.py

@@ -1,96 +1,52 @@
+from torch.utils.data import Dataset
+from transformers import BertTokenizer, BertForMaskedLM
 import torch
-from torch.utils.data import DataLoader, Dataset
-from transformers import AutoTokenizer, AutoModelForMaskedLM, pipeline
-from opacus import PrivacyEngine
-import gradio as gr
 
-# ===========================
-# مدل و توکنایزر
-# ===========================
-model_name = "HooshvareLab/bert-base-parsbert-uncased"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-tokenizer.pad_token = tokenizer.eos_token  # حل ارور padding
+# بارگذاری توکنایزر و مدل
+tokenizer = BertTokenizer.from_pretrained("HooshvareLab/bert-base-parsbert-uncased")
+model = BertForMaskedLM.from_pretrained("HooshvareLab/bert-base-parsbert-uncased")
 
-model = AutoModelForMaskedLM.from_pretrained(model_name)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+# اضافه کردن pad_token در صورت نیاز
+if tokenizer.pad_token is None:
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    model.resize_token_embeddings(len(tokenizer))
 
-# ===========================
-# بارگذاری داده آموزشی داخلی
-# ===========================
-with open("training_data.txt", "r", encoding="utf-8") as f:
-    texts = [line.strip() for line in f if line.strip()]
-
-# ===========================
-# Dataset و DataLoader
-# ===========================
 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 self.encodings.input_ids.size(0)
-    def __getitem__(self, idx):
-        return {"input_ids": self.encodings.input_ids[idx], "attention_mask": self.encodings.attention_mask[idx], "labels": self.encodings.input_ids[idx]}
-
-dataset = TextDataset(texts, tokenizer)
-data_loader = DataLoader(dataset, batch_size=4, shuffle=True)
-
-# ===========================
-# تعریف Optimizer و PrivacyEngine
-# ===========================
-optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
-
-privacy_engine = PrivacyEngine()
-model, optimizer, data_loader = privacy_engine.make_private(
-    module=model,
-    optimizer=optimizer,
-    data_loader=data_loader,
-    max_grad_norm=1.0,
-    noise_multiplier=1.0,
-)
+    def __init__(self, file_path, tokenizer, max_length=128):
+        # خواندن خطوط غیرخالی از فایل
+        with open(file_path, 'r', encoding='utf-8') as f:
+            self.texts = [line.strip() for line in f if line.strip()]
+        self.tokenizer = tokenizer
+        self.max_length = max_length
 
-# ===========================
-# آموزش incremental (یک epoch ساده)
-# ===========================
-def train_step():
-    model.train()
-    total_loss = 0
-    for batch in data_loader:
-        input_ids = batch["input_ids"].to(device)
-        attention_mask = batch["attention_mask"].to(device)
-        labels = batch["labels"].to(device)
-
-        optimizer.zero_grad()
-        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
-        loss = outputs.loss
-        loss.backward()
-        optimizer.step()
-        total_loss += loss.item()
-    return total_loss / len(data_loader)
-
-# ===========================
-# تابع تولید پاسخ
-# ===========================
-generator = pipeline(
-    task="text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0 if torch.cuda.is_available() else -1
-)
-
-def chat(prompt):
-    loss = train_step()  # آموزش incremental روی داده داخلی قبل از تولید پاسخ
-    result = generator(prompt, max_length=128, do_sample=True, temperature=0.8)
-    return f"{result[0]['generated_text']}\n\n(مدل به صورت incremental روی داده‌های داخلی آموزش داده شد. DP فعال است)"
-
-# ===========================
-# رابط کاربری Gradio (RTL)
-# ===========================
-with gr.Blocks() as demo:
-    gr.HTML("<div dir='rtl'><h2>چت بات فارسی با یادگیری امن</h2></div>")
-    prompt_input = gr.Textbox(label="متن خود را وارد کنید:", placeholder="اینجا بنویسید...", lines=3)
-    output = gr.Textbox(label="پاسخ مدل:", lines=8)
-    btn = gr.Button("ارسال")
-    btn.click(chat, inputs=prompt_input, outputs=output)
+    def __len__(self):
+        return len(self.texts)
 
-demo.launch()
+    def __getitem__(self, idx):
+        text = self.texts[idx]
+        encoding = self.tokenizer(
+            text,
+            truncation=True,
+            padding='max_length',
+            max_length=self.max_length,
+            return_tensors='pt'
+        )
+        input_ids = encoding['input_ids'].squeeze(0)
+        attention_mask = encoding['attention_mask'].squeeze(0)
+        
+        # برای MaskedLM، label همان input_ids است، ولی padding با -100 نشانه‌گذاری می‌شود تا در loss حساب نشود
+        labels = input_ids.clone()
+        labels[input_ids == tokenizer.pad_token_id] = -100
+        
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'labels': labels
+        }
+
+# نمونه‌سازی Dataset با فایل متنی آپلود شده
+dataset = TextDataset("training.txt", tokenizer)
+
+# بررسی دو نمونه
+for i in range(2):
+    print(dataset[i])