Update app.py

app.py

@@ -1,10 +1,7 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import numpy as np
 import gradio as gr
-from opacus import PrivacyEngine
-from torch.utils.data import Dataset
-from torch.optim import AdamW
 from sentence_transformers import SentenceTransformer
 import faiss
 
@@ -33,20 +30,6 @@ embedder = None
 index = None
 texts = []
 
-# Custom Dataset for training data
-class ChatDataset(Dataset):
-    def __init__(self, data):
-        self.data = data
-        self.encodings = tokenizer(data, truncation=True, padding=True, max_length=128, return_tensors="pt")
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        item = {key: val[idx].to(device) for key, val in self.encodings.items()}
-        item["labels"] = item["input_ids"].clone()
-        return item
-
 # Load training data from training_data.txt in the root directory
 def load_training_data():
     global texts
@@ -72,63 +55,18 @@ def build_rag_index(texts):
     index.add(embeddings)
     return embedder, index
 
-# Fine-tune model with differential privacy
+# Fine-tune model with differential privacy (skipped to use only pretrained LLaMA 2)
 def train_model():
     global texts, embedder, index
     texts = load_training_data()
     if not texts:
-        print("No training data available. Skipping training.")
+        print("No training data available. Skipping RAG index build.")
         return
 
     # Build RAG index
     build_rag_index(texts)
 
-    train_dataset = ChatDataset(texts)
-
-    training_args = TrainingArguments(
-        output_dir="./results",
-        num_train_epochs=3,
-        per_device_train_batch_size=4,  # Reduced to avoid memory issues
-        save_steps=10_000,
-        save_total_limit=2,
-        use_cpu=True if not torch.cuda.is_available() else False,  # Use CPU if no GPU
-        logging_steps=100,  # Log training progress
-    )
-
-    # Initialize optimizer explicitly
-    optimizer = AdamW(model.parameters(), lr=5e-5)
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        optimizers=(optimizer, None),  # Pass optimizer explicitly
-    )
-
-    # Set model to training mode for differential privacy
-    model.train()
-
-    # Add differential privacy
-    try:
-        privacy_engine = PrivacyEngine(secure_mode=False)  # False for experimentation
-        private_model, private_optimizer, train_dataloader = privacy_engine.make_private(
-            module=model,
-            optimizer=optimizer,
-            data_loader=trainer.get_train_dataloader(),
-            noise_multiplier=1.1,
-            max_grad_norm=1.0,
-        )
-
-        trainer.optimizer = private_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")
-    except Exception as e:
-        print(f"Error during training with differential privacy: {e}")
+    print("Skipping fine-tuning to use only pretrained LLaMA 2 model.")
 
 def add_noise(tensor, sensitivity, epsilon, delta):
     """Add Laplace noise for differential privacy."""
@@ -180,14 +118,14 @@ def chat(message, history):
     
     return response
 
-# Train the model on startup
+# Train the model on startup (now only loads data and builds RAG index)
 train_model()
 
 # Gradio interface
 iface = gr.ChatInterface(
     fn=chat,
-    title="LLaMA 2 Persian Chatbot with Differential Privacy and RAG",
-    description="Chat with a fine-tuned LLaMA 2 Persian model using training_data.txt from the root directory with RAG."
+    title="LLaMA 2 Persian Chatbot with RAG",
+    description="Chat with pretrained LLaMA 2 Persian model using training_data.txt as RAG knowledge base."
 )
 
 if __name__ == "__main__":