Update app.py

app.py

@@ -1,17 +1,14 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
-from datasets import load_dataset
+from transformers import BertLMHeadModel, AutoTokenizer, TrainingArguments, Trainer
 import numpy as np
-from flask import Flask, request, jsonify
+import gradio as gr
 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)
+model = BertLMHeadModel.from_pretrained(model_name, is_decoder=True)
 
 # Differential Privacy parameters
 epsilon = 1.0  # Privacy budget
@@ -25,7 +22,7 @@ conversation_history = []
 class ChatDataset(Dataset):
     def __init__(self, data):
         self.data = data
-        self.encodings = tokenizer(data, truncation=True, padding=True, return_tensors="pt")
+        self.encodings = tokenizer(data, truncation=True, padding=True, max_length=128, return_tensors="pt")
 
     def __len__(self):
         return len(self.data)
@@ -35,11 +32,12 @@ class ChatDataset(Dataset):
         item["labels"] = item["input_ids"].clone()
         return item
 
-# Load training data from Hugging Face
+# Load training data from training_data.txt in the root directory
 def load_training_data():
     try:
         with open("training_data.txt", "r", encoding="utf-8") as file:
             texts = [line.strip() for line in file if line.strip()]
+        print(f"Loaded {len(texts)} training examples from training_data.txt")
         return texts
     except FileNotFoundError:
         print("Error: training_data.txt not found in the root directory.")
@@ -51,14 +49,20 @@ def load_training_data():
 # Fine-tune model with differential privacy
 def train_model():
     texts = load_training_data()
+    if not texts:
+        print("No training data available. Skipping training.")
+        return
+
     train_dataset = ChatDataset(texts)
 
     training_args = TrainingArguments(
         output_dir="./results",
         num_train_epochs=3,
-        per_device_train_batch_size=8,
+        per_device_train_batch_size=4,  # Reduced to avoid memory issues
         save_steps=10_000,
         save_total_limit=2,
+        no_cuda=True if not torch.cuda.is_available() else None,  # CPU compatibility
+        logging_steps=100,  # Log training progress
     )
 
     trainer = Trainer(
@@ -69,7 +73,7 @@ def train_model():
 
     # Add differential privacy
     privacy_engine = PrivacyEngine()
-    model, optimizer, train_dataloader = privacy_engine.make_private(
+    private_model, optimizer, train_dataloader = privacy_engine.make_private(
         module=model,
         optimizer=trainer.optimizer,
         data_loader=trainer.get_train_dataloader(),
@@ -79,16 +83,18 @@ def train_model():
 
     trainer.optimizer = optimizer
     trainer.train_dataloader = train_dataloader
+    trainer.model = private_model  # Update trainer to use private model
 
     trainer.train()
     model.save_pretrained("./fine_tuned_model")
     tokenizer.save_pretrained("./fine_tuned_model")
+    print("Model training completed and saved to ./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)
+    return tensor + torch.tensor(noise, dtype=tensor.dtype, device=tensor.device)
 
 def update_model(user_input, response):
     """Update conversation history."""
@@ -96,34 +102,38 @@ def update_model(user_input, response):
     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}")
+    return f"Learning from: {user_input} -> {response}"
 
-@app.route('/chat', methods=['POST'])
-def chat():
-    data = request.json
-    user_input = data.get('message', '')
-    
+def chat(message, history):
     # Tokenize input
-    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
+    inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=128).to(model.device)
     
-    # Generate response with model
+    # Generate response with model using beam search
     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)
+        outputs = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            max_length=50,
+            num_beams=5,
+            no_repeat_ngram_size=2,
+            early_stopping=True,
+        )
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     # Update conversation history
-    update_model(user_input, response)
+    update_model(message, response)
     
-    return jsonify({"response": response})
+    return response
+
+# Train the model on startup
+train_model()
+
+# Gradio interface
+iface = gr.ChatInterface(
+    fn=chat,
+    title="ParsBERT Chatbot with Differential Privacy",
+    description="Chat with a fine-tuned ParsBERT model using training_data.txt from the root directory."
+)
 
-if __name__ == '__main__':
-    # Train the model on startup
-    train_model()
-    app.run(debug=True)
+if __name__ == "__main__":
+    iface.launch(server_name="0.0.0.0", server_port=7860)