Update app.py

app.py

@@ -12,8 +12,8 @@ torch.set_default_dtype(torch.float32)
 # Set device explicitly
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load LLaMA 3.2 1B model (no quantization for CPU compatibility)
-model_name = "meta-llama/Llama-3.2-1B-Instruct"
+# Load Persian GPT-2 model and tokenizer
+model_name = "HooshvareLab/gpt2-fa"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -75,7 +75,7 @@ def train_model():
 
     # Build RAG index
     build_rag_index(texts)
-    print("Using pretrained LLaMA 3.2 1B model without fine-tuning.")
+    print("Using pretrained Persian GPT-2 model without fine-tuning.")
 
 def add_noise(tensor, sensitivity, epsilon, delta):
     """Add Laplace noise for differential privacy."""
@@ -100,7 +100,7 @@ def chat(message, history):
     if embedder and index:
         try:
             query_emb = embedder.encode(message, convert_to_tensor=True).cpu().numpy()
-            D, I = index.search(query_emb, k=3)
+            D, I = index.search(query_emb, k=5)  # Increased k for better context
             retrieved = [texts[i] for i in I[0] if i >= 0 and i < len(texts)]
             context = "\n".join(retrieved)
         except Exception as e:
@@ -118,7 +118,7 @@ def chat(message, history):
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
             max_length=100,  # Increased for better responses
-            num_beams=5,
+            num_beams=7,    # Increased for better quality
             no_repeat_ngram_size=2,
             early_stopping=True,
         )
@@ -142,8 +142,8 @@ train_model()
 # Gradio interface
 iface = gr.ChatInterface(
     fn=chat,
-    title="LLaMA 3.2 1B Persian Chatbot with RAG",
-    description="Chat with pretrained LLaMA 3.2 1B model using training_data.txt as RAG knowledge base."
+    title="Persian GPT-2 Chatbot with RAG",
+    description="Chat with pretrained Persian GPT-2 model using training_data.txt as RAG knowledge base."
 )
 
 if __name__ == "__main__":