Update app.py

app.py

@@ -1,10 +1,12 @@
 import torch
-from transformers import BertLMHeadModel, AutoTokenizer, TrainingArguments, Trainer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 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
 
 # Disable torch.compile to avoid meta device issues
 torch._dynamo.config.suppress_errors = True
@@ -13,10 +15,10 @@ torch.set_default_dtype(torch.float32)
 # Set device explicitly
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load ParsBERT model and tokenizer
-model_name = "HooshvareLab/bert-base-parsbert-uncased"
+# Load LLaMA 2 Persian model and tokenizer
+model_name = "sinarashidi/llama-2-7b-chat-persian"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = BertLMHeadModel.from_pretrained(model_name, is_decoder=True).to(device)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32).to(device)
 
 # Differential Privacy parameters
 epsilon = 1.0  # Privacy budget
@@ -26,6 +28,11 @@ sensitivity = 1.0  # Sensitivity of the query
 # Simple memory for conversation history
 conversation_history = []
 
+# RAG components
+embedder = None
+index = None
+texts = []
+
 # Custom Dataset for training data
 class ChatDataset(Dataset):
     def __init__(self, data):
@@ -42,6 +49,7 @@ class ChatDataset(Dataset):
 
 # Load training data from training_data.txt in the root directory
 def load_training_data():
+    global texts
     try:
         with open("training_data.txt", "r", encoding="utf-8") as file:
             texts = [line.strip() for line in file if line.strip()]
@@ -54,13 +62,27 @@ def load_training_data():
         print(f"Error reading training_data.txt: {e}")
         return []
 
+# Build RAG index
+def build_rag_index(texts):
+    global embedder, index
+    embedder = SentenceTransformer('xmanii/maux-gte-persian')
+    embeddings = embedder.encode(texts, convert_to_tensor=True).cpu().numpy()
+    dimension = embeddings.shape[1]
+    index = faiss.IndexFlatL2(dimension)
+    index.add(embeddings)
+    return embedder, index
+
 # Fine-tune model with differential privacy
 def train_model():
+    global texts, embedder, index
     texts = load_training_data()
     if not texts:
         print("No training data available. Skipping training.")
         return
 
+    # Build RAG index
+    build_rag_index(texts)
+
     train_dataset = ChatDataset(texts)
 
     training_args = TrainingArguments(
@@ -126,8 +148,20 @@ def chat(message, history):
     # Set model to evaluation mode for inference
     model.eval()
     
+    # RAG retrieval
+    if embedder and index:
+        query_emb = embedder.encode(message, convert_to_tensor=True).cpu().numpy()
+        D, I = index.search(query_emb, k=3)
+        retrieved = [texts[i] for i in I[0] if i >= 0 and i < len(texts)]
+        context = "\n".join(retrieved)
+    else:
+        context = ""
+
+    # Prepare prompt with context
+    prompt = f"Context: {context}\nUser: {message}\nBot:"
+    
     # Tokenize input
-    inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
     
     # Generate response with model using beam search
     with torch.no_grad():
@@ -152,8 +186,8 @@ 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."
+    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."
 )
 
 if __name__ == "__main__":