Update app.py

app.py

@@ -1,52 +1,122 @@
-from torch.utils.data import Dataset
-from transformers import BertTokenizer, BertForMaskedLM
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from datasets import load_dataset
+import numpy as np
+from flask import Flask, request, jsonify
+from opacus import PrivacyEngine
+from torch.utils.data import Dataset
+
+app = Flask(__name__)
+
+# Load ParsBERT model and tokenizer
+model_name = "HooshvareLab/bert-base-parsbert-uncased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# بارگذاری توکنایزر و مدل
-tokenizer = BertTokenizer.from_pretrained("HooshvareLab/bert-base-parsbert-uncased")
-model = BertForMaskedLM.from_pretrained("HooshvareLab/bert-base-parsbert-uncased")
+# Differential Privacy parameters
+epsilon = 1.0  # Privacy budget
+delta = 1e-5   # Privacy parameter
+sensitivity = 1.0  # Sensitivity of the query
 
-# اضافه کردن pad_token در صورت نیاز
-if tokenizer.pad_token is None:
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    model.resize_token_embeddings(len(tokenizer))
+# Simple memory for conversation history
+conversation_history = []
 
-class TextDataset(Dataset):
-    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
+# Custom Dataset for training data
+class ChatDataset(Dataset):
+    def __init__(self, data):
+        self.data = data
+        self.encodings = tokenizer(data, truncation=True, padding=True, return_tensors="pt")
 
     def __len__(self):
-        return len(self.texts)
+        return len(self.data)
 
     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_data.txt", tokenizer)
-
-# بررسی دو نمونه
-for i in range(2):
-    print(dataset[i])
+        item = {key: val[idx] for key, val in self.encodings.items()}
+        item["labels"] = item["input_ids"].clone()
+        return item
+
+# Load training data from Hugging Face
+def load_training_data():
+    dataset = load_dataset("training_data.txt", split="train")  # Replace with actual dataset path
+    texts = [item["text"] for item in dataset if "text" in item]
+    return texts
+
+# Fine-tune model with differential privacy
+def train_model():
+    texts = load_training_data()
+    train_dataset = ChatDataset(texts)
+
+    training_args = TrainingArguments(
+        output_dir="./results",
+        num_train_epochs=3,
+        per_device_train_batch_size=8,
+        save_steps=10_000,
+        save_total_limit=2,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+    )
+
+    # Add differential privacy
+    privacy_engine = PrivacyEngine()
+    model, optimizer, train_dataloader = privacy_engine.make_private(
+        module=model,
+        optimizer=trainer.optimizer,
+        data_loader=trainer.get_train_dataloader(),
+        noise_multiplier=1.1,
+        max_grad_norm=1.0,
+    )
+
+    trainer.optimizer = optimizer
+    trainer.train_dataloader = train_dataloader
+
+    trainer.train()
+    model.save_pretrained("./fine_tuned_model")
+    tokenizer.save_pretrained("./fine_tuned_model")
+
+def add_noise(tensor, sensitivity, epsilon, delta):
+    """Add Laplace noise for differential privacy."""
+    scale = sensitivity / epsilon
+    noise = np.random.laplace(0, scale, tensor.shape)
+    return tensor + torch.tensor(noise, dtype=tensor.dtype)
+
+def update_model(user_input, response):
+    """Update conversation history."""
+    global conversation_history
+    conversation_history.append({"user": user_input, "bot": response})
+    if len(conversation_history) > 100:  # Limit history size
+        conversation_history.pop(0)
+    print(f"Learning from: {user_input} -> {response}")
+
+@app.route('/chat', methods=['POST'])
+def chat():
+    data = request.json
+    user_input = data.get('message', '')
+    
+    # Tokenize input
+    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
+    
+    # Generate response with model
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+    
+    # Apply differential privacy to logits
+    noisy_logits = add_noise(logits, sensitivity, epsilon, delta)
+    
+    # Decode response
+    response_ids = torch.argmax(noisy_logits, dim=-1)
+    response = tokenizer.decode(response_ids[0], skip_special_tokens=True)
+    
+    # Update conversation history
+    update_model(user_input, response)
+    
+    return jsonify({"response": response})
+
+if __name__ == '__main__':
+    # Train the model on startup
+    train_model()
+    app.run(debug=True)