Update app.py

app.py

@@ -16,19 +16,13 @@ from transformers import (
     BitsAndBytesConfig,
 )
 
-# PEFT (LoRA / QLoRA)
 from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training, PeftModel
 
-# For embeddings
 from sentence_transformers import SentenceTransformer
 
-##############################################################################
-# QLoRA Demo Setup
-##############################################################################
-
 TEXT_PIPELINE = None
 COMPARISON_PIPELINE = None
-NUM_EXAMPLES = 50  # We'll train on 50 rows for demonstration
+NUM_EXAMPLES = 50  
 
 @spaces.GPU(duration=300)
 def finetune_small_subset():
@@ -40,7 +34,6 @@ def finetune_small_subset():
     5) Reloads LoRA adapters for inference in a pipeline.
     """
 
-    # --- 1) Load a small subset of the Magpie dataset ---
     ds = load_dataset(
         "Magpie-Align/Magpie-Reasoning-V2-250K-CoT-Deepseek-R1-Llama-70B", 
         split="train"
@@ -52,10 +45,9 @@ def finetune_small_subset():
 
     ds = ds.select(range(min(NUM_EXAMPLES, len(ds))))
 
-    # --- 2) Setup 4-bit quantization ---
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,  # or torch.float16
+        bnb_4bit_compute_dtype=torch.bfloat16,  
         bnb_4bit_use_double_quant=True,
         bnb_4bit_quant_type="nf4",
     )
@@ -75,14 +67,13 @@ def finetune_small_subset():
         "wuhp/myr1",
         subfolder="myr1",
         config=config,
-        quantization_config=bnb_config,  # <--- QLoRA 4-bit
+        quantization_config=bnb_config,  
         device_map="auto",
         trust_remote_code=True
     )
 
     base_model = prepare_model_for_kbit_training(base_model)
 
-    # --- 3) Create LoRA config & wrap the base model in LoRA ---
     lora_config = LoraConfig(
         r=16,
         lora_alpha=32,
@@ -93,7 +84,6 @@ def finetune_small_subset():
     )
     lora_model = get_peft_model(base_model, lora_config)
 
-    # --- 4) Tokenize dataset ---
     def tokenize_fn(ex):
         text = (
             f"Instruction: {ex['instruction']}\n\n"
@@ -106,7 +96,6 @@ def finetune_small_subset():
 
     collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-    # Training args
     training_args = TrainingArguments(
         output_dir="finetuned_myr1",
         num_train_epochs=1,
@@ -126,11 +115,9 @@ def finetune_small_subset():
     )
     trainer.train()
 
-    # --- 5) Save LoRA adapter + tokenizer ---
     trainer.model.save_pretrained("finetuned_myr1")
     tokenizer.save_pretrained("finetuned_myr1")
 
-    # --- 6) Reload for inference
     base_model_2 = AutoModelForCausalLM.from_pretrained(
         "wuhp/myr1",
         subfolder="myr1",
@@ -235,9 +222,6 @@ def compare_models(prompt, temperature, top_p, min_new_tokens, max_new_tokens):
     )
     return local_out[0]["generated_text"], comp_out[0]["generated_text"]
 
-###############################################################################
-# Retrieval-Augmented Memory with FAISS
-###############################################################################
 class ConversationRetriever:
     """
     A simple in-memory store + FAISS for retrieval of conversation chunks.
@@ -253,12 +237,10 @@ class ConversationRetriever:
         self.embed_model = SentenceTransformer(model_name)
         self.embed_dim = embed_dim
 
-        # We'll store (text, vector) in FAISS.  For metadata, store in python list/dict.
-        # For a real app, you'd probably want a more robust store.
         self.index = faiss.IndexFlatL2(embed_dim)
-        self.texts = []  # store the raw text chunks
-        self.vectors = []  # store vectors (redundant but simpler to show)
-        self.ids = []  # store an integer ID or similar
+        self.texts = []  
+        self.vectors = []  
+        self.ids = []  
 
         self.id_counter = 0
 
@@ -271,7 +253,7 @@ class ConversationRetriever:
             return
 
         emb = self.embed_model.encode([text], convert_to_numpy=True)
-        vec = emb[0].astype(np.float32)  # shape [embed_dim]
+        vec = emb[0].astype(np.float32)  
         self.index.add(vec.reshape(1, -1))
 
         self.texts.append(text)
@@ -288,17 +270,13 @@ class ConversationRetriever:
         q_vec = q_emb[0].reshape(1, -1)
         distances, indices = self.index.search(q_vec, top_k)
 
-        # indices is shape [1, top_k], distances is shape [1, top_k]
         results = []
         for dist, idx in zip(distances[0], indices[0]):
-            if idx < len(self.texts):  # safety check
+            if idx < len(self.texts):  
                 results.append((self.texts[idx], dist))
         return results
 
-###############################################################################
-# Build a Chat that uses RAG
-###############################################################################
-retriever = ConversationRetriever()  # global retriever instance
+retriever = ConversationRetriever()  
 
 def build_rag_prompt(user_query, retrieved_chunks):
     """
@@ -331,18 +309,13 @@ def chat_rag(user_input, history, temperature, top_p, min_new_tokens, max_new_to
     """
     pipe = ensure_pipeline()
 
-    # 1) Add the user input as a chunk to the retriever DB.
     retriever.add_text(f"User: {user_input}")
 
-    # 2) Retrieve top-3 older chunks. We can skip the chunk we just added if we want to
-    # (since it's the same text), but for simplicity let's just do a search for user_input.
     top_k = 3
     results = retriever.search(user_input, top_k=top_k)
 
-    # 3) Build final prompt
     prompt = build_rag_prompt(user_input, results)
 
-    # 4) Generate
     output = pipe(
         prompt,
         temperature=float(temperature),
@@ -352,23 +325,16 @@ def chat_rag(user_input, history, temperature, top_p, min_new_tokens, max_new_to
         do_sample=True
     )[0]["generated_text"]
 
-    # We only want the new part after "Assistant:"
-    # Because the pipeline output includes the entire prompt + new text.
     if output.startswith(prompt):
         assistant_reply = output[len(prompt):].strip()
     else:
         assistant_reply = output.strip()
 
-    # 5) Add the assistant's response to the DB as well
     retriever.add_text(f"Assistant: {assistant_reply}")
 
-    # 6) Update the chat history for display in the Gradio Chatbot
     history.append([user_input, assistant_reply])
     return history, history
 
-###############################################################################
-# Gradio UI
-###############################################################################
 with gr.Blocks() as demo:
     gr.Markdown("# QLoRA Fine-tuning & RAG-based Chat Demo")
 
@@ -377,7 +343,6 @@ with gr.Blocks() as demo:
 
     finetune_btn.click(fn=finetune_small_subset, outputs=status_box)
 
-    # Simple generation UI (no retrieval):
     gr.Markdown("## Direct Generation (No Retrieval)")
     prompt_in = gr.Textbox(lines=3, label="Prompt")
     temperature = gr.Slider(0.0, 1.5, step=0.1, value=0.7, label="Temperature")
@@ -393,7 +358,6 @@ with gr.Blocks() as demo:
         outputs=output_box
     )
 
-    # Comparison UI:
     gr.Markdown("## Compare myr1 vs DeepSeek")
     compare_btn = gr.Button("Compare")
     out_local = gr.Textbox(label="myr1 Output", lines=6)
@@ -404,12 +368,11 @@ with gr.Blocks() as demo:
         outputs=[out_local, out_deepseek]
     )
 
-    # RAG-based Chat
     gr.Markdown("## Chat with Retrieval-Augmented Memory")
     with gr.Row():
         with gr.Column():
             chatbot = gr.Chatbot(label="RAG Chat")
-            chat_state = gr.State([])  # just for display
+            chat_state = gr.State([])  
 
             user_input = gr.Textbox(
                 show_label=False,
@@ -418,7 +381,6 @@ with gr.Blocks() as demo:
             )
             send_btn = gr.Button("Send")
 
-    # On user submit, call chat_rag
     user_input.submit(
         fn=chat_rag,
         inputs=[user_input, chat_state, temperature, top_p, min_tokens, max_tokens],