Update app.py

app.py

@@ -20,6 +20,11 @@ from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_tr
 
 from sentence_transformers import SentenceTransformer
 
+# Import your custom configuration overrides.
+# For example, your configuration_deepseek.py might export a dictionary called CONFIG_OVERRIDES.
+import configuration_deepseek
+
+# Global variables for pipelines and settings.
 TEXT_PIPELINE = None
 COMPARISON_PIPELINE = None
 NUM_EXAMPLES = 1000  
@@ -27,47 +32,46 @@ NUM_EXAMPLES = 1000
 @spaces.GPU(duration=300)
 def finetune_small_subset():
     """
-    1) Loads 'wuhp/myr1' in 4-bit quantization (QLoRA style),
+    1) Loads your custom model ("wuhp/myr1") in 4-bit quantization (QLoRA style),
     2) Adds LoRA adapters (trainable),
-    3) Trains on a small subset of the Magpie dataset,
-    4) Saves LoRA adapter to 'finetuned_myr1',
-    5) Reloads LoRA adapters for inference in a pipeline.
+    3) Fine-tunes on a small subset of the ServiceNow-AI/R1-Distill-SFT dataset,
+    4) Saves the LoRA adapter to "finetuned_myr1",
+    5) Reloads the LoRA adapter for inference.
     """
-
-    ds = load_dataset(
-        "Magpie-Align/Magpie-Reasoning-V2-250K-CoT-Deepseek-R1-Llama-70B", 
-        split="train"
-    )
-
-    unique_ids = list(set(ds["conversation_id"]))
-    single_id = unique_ids[0]
-    ds = ds.filter(lambda x: x["conversation_id"] == single_id)
-
+    # Load the new dataset.
+    ds = load_dataset("ServiceNow-AI/R1-Distill-SFT", split="train")
     ds = ds.select(range(min(NUM_EXAMPLES, len(ds))))
 
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,  
+        bnb_4bit_compute_dtype=torch.bfloat16,
         bnb_4bit_use_double_quant=True,
         bnb_4bit_quant_type="nf4",
     )
 
-    config = AutoConfig.from_pretrained(
-        "wuhp/myr1", 
+    # Load the base configuration from your model repository.
+    base_config = AutoConfig.from_pretrained(
+        "wuhp/myr1",
         subfolder="myr1",
-        trust_remote_code=True
+        trust_remote_code=True,
     )
+    # Apply your custom overrides (from configuration_deepseek.py).
+    for key, value in configuration_deepseek.CONFIG_OVERRIDES.items():
+        setattr(base_config, key, value)
+
     tokenizer = AutoTokenizer.from_pretrained(
-        "wuhp/myr1", 
+        "wuhp/myr1",
         subfolder="myr1",
         trust_remote_code=True
     )
 
+    # Load the model. With trust_remote_code=True, your custom model class (e.g. DeepseekV3ForCausalLM)
+    # will be loaded from the repository.
     base_model = AutoModelForCausalLM.from_pretrained(
         "wuhp/myr1",
         subfolder="myr1",
-        config=config,
-        quantization_config=bnb_config,  
+        config=base_config,
+        quantization_config=bnb_config,
         device_map="auto",
         trust_remote_code=True
     )
@@ -84,10 +88,11 @@ def finetune_small_subset():
     )
     lora_model = get_peft_model(base_model, lora_config)
 
+    # For this dataset, assume "problem" is the prompt and "solution" is the target.
     def tokenize_fn(ex):
         text = (
-            f"Instruction: {ex['instruction']}\n\n"
-            f"Response: {ex['response']}"
+            f"Problem: {ex['problem']}\n\n"
+            f"Solution: {ex['solution']}"
         )
         return tokenizer(text, truncation=True, max_length=512)
 
@@ -102,9 +107,9 @@ def finetune_small_subset():
         per_device_train_batch_size=1,
         gradient_accumulation_steps=2,
         logging_steps=5,
-        save_steps=999999,   
+        save_steps=999999,   # High save interval
         save_total_limit=1,
-        fp16=False,          
+        fp16=False,          # Set to True if supported by your hardware
     )
 
     trainer = Trainer(
@@ -115,13 +120,15 @@ def finetune_small_subset():
     )
     trainer.train()
 
+    # Save the LoRA adapter and tokenizer.
     trainer.model.save_pretrained("finetuned_myr1")
     tokenizer.save_pretrained("finetuned_myr1")
 
+    # Reload the base model and attach the LoRA adapter for inference.
     base_model_2 = AutoModelForCausalLM.from_pretrained(
         "wuhp/myr1",
         subfolder="myr1",
-        config=config,
+        config=base_config,
         quantization_config=bnb_config,
         device_map="auto",
         trust_remote_code=True
@@ -140,8 +147,8 @@ def finetune_small_subset():
 
 def ensure_pipeline():
     """
-    If we haven't finetuned yet (TEXT_PIPELINE is None),
-    load the base model in 4-bit with NO LoRA.
+    If we haven't fine-tuned yet (i.e. TEXT_PIPELINE is None),
+    load the base model (without LoRA) in 4-bit mode.
     """
     global TEXT_PIPELINE
     if TEXT_PIPELINE is None:
@@ -151,12 +158,14 @@ def ensure_pipeline():
             bnb_4bit_use_double_quant=True,
             bnb_4bit_quant_type="nf4",
         )
-        config = AutoConfig.from_pretrained("wuhp/myr1", subfolder="myr1", trust_remote_code=True)
+        base_config = AutoConfig.from_pretrained("wuhp/myr1", subfolder="myr1", trust_remote_code=True)
+        for key, value in configuration_deepseek.CONFIG_OVERRIDES.items():
+            setattr(base_config, key, value)
         tokenizer = AutoTokenizer.from_pretrained("wuhp/myr1", subfolder="myr1", trust_remote_code=True)
         base_model = AutoModelForCausalLM.from_pretrained(
             "wuhp/myr1",
             subfolder="myr1",
-            config=config,
+            config=base_config,
             quantization_config=bnb_config,
             device_map="auto",
             trust_remote_code=True
@@ -166,7 +175,7 @@ def ensure_pipeline():
 
 def ensure_comparison_pipeline():
     """
-    Load the DeepSeek model pipeline if not already loaded.
+    Load a reference DeepSeek model pipeline if not already loaded.
     """
     global COMPARISON_PIPELINE
     if COMPARISON_PIPELINE is None:
@@ -183,7 +192,7 @@ def ensure_comparison_pipeline():
 @spaces.GPU(duration=120)
 def predict(prompt, temperature, top_p, min_new_tokens, max_new_tokens):
     """
-    Simple single-prompt generation (no retrieval).
+    Direct generation without retrieval.
     """
     pipe = ensure_pipeline()
     out = pipe(
@@ -199,7 +208,7 @@ def predict(prompt, temperature, top_p, min_new_tokens, max_new_tokens):
 @spaces.GPU(duration=120)
 def compare_models(prompt, temperature, top_p, min_new_tokens, max_new_tokens):
     """
-    Compare local pipeline vs. DeepSeek side-by-side.
+    Compare outputs between your custom model and a reference DeepSeek model.
     """
     local_pipe = ensure_pipeline()
     comp_pipe = ensure_comparison_pipeline()
@@ -224,75 +233,51 @@ def compare_models(prompt, temperature, top_p, min_new_tokens, max_new_tokens):
 
 class ConversationRetriever:
     """
-    A simple in-memory store + FAISS for retrieval of conversation chunks.
-    Each chunk is embedded via SentenceTransformer. On a new user query,
-    we embed the query, do similarity search, and retrieve top-k relevant chunks.
+    A simple in-memory FAISS-based retriever.
+    Each text chunk is embedded using SentenceTransformer.
     """
-
     def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2", embed_dim=384):
-        """
-        model_name: embedding model for messages
-        embed_dim: dimension of the embeddings from that model
-        """
         self.embed_model = SentenceTransformer(model_name)
         self.embed_dim = embed_dim
-
         self.index = faiss.IndexFlatL2(embed_dim)
-        self.texts = []  
-        self.vectors = []  
-        self.ids = []  
-
+        self.texts = []
+        self.vectors = []
+        self.ids = []
         self.id_counter = 0
 
     def add_text(self, text):
-        """
-        Add a new text chunk to the vector store.
-        Could chunk it up if desired, but here we treat the entire text as one chunk.
-        """
         if not text.strip():
             return
-
         emb = self.embed_model.encode([text], convert_to_numpy=True)
-        vec = emb[0].astype(np.float32)  
+        vec = emb[0].astype(np.float32)
         self.index.add(vec.reshape(1, -1))
-
         self.texts.append(text)
         self.vectors.append(vec)
         self.ids.append(self.id_counter)
-
         self.id_counter += 1
 
     def search(self, query, top_k=3):
-        """
-        Given a query, embed it, do similarity search in FAISS, return top-k texts.
-        """
         q_emb = self.embed_model.encode([query], convert_to_numpy=True).astype(np.float32)
         q_vec = q_emb[0].reshape(1, -1)
         distances, indices = self.index.search(q_vec, top_k)
-
         results = []
         for dist, idx in zip(distances[0], indices[0]):
-            if idx < len(self.texts):  
+            if idx < len(self.texts):
                 results.append((self.texts[idx], dist))
         return results
 
-retriever = ConversationRetriever()  
+retriever = ConversationRetriever()
 
 def build_rag_prompt(user_query, retrieved_chunks):
     """
-    Construct a prompt that includes:
-      - The user's new query
-      - A "Relevant Context" section from retrieved chunks
-      - "Assistant:" to let the model continue
-    Feel free to customize the formatting as you like.
+    Build a prompt for retrieval-augmented generation.
     """
     context_str = ""
     for i, (chunk, dist) in enumerate(retrieved_chunks):
-        context_str += f"Chunk #{i+1} (similarity score ~ {dist:.2f}):\n{chunk}\n\n"
-
+        context_str += f"Chunk #{i+1} (similarity ~ {dist:.2f}):\n{chunk}\n\n"
     prompt = (
         f"User's Query:\n{user_query}\n\n"
-        f"Relevant Context from Conversation:\n{context_str}"
+        f"Relevant Context:\n{context_str}"
         "Assistant:"
     )
     return prompt
@@ -300,22 +285,13 @@ def build_rag_prompt(user_query, retrieved_chunks):
 @spaces.GPU(duration=120)
 def chat_rag(user_input, history, temperature, top_p, min_new_tokens, max_new_tokens):
     """
-    Our RAG-based chat function. We'll:
-      1) Add user input to FAISS
-      2) Retrieve top-k relevant older messages from FAISS
-      3) Build a prompt that includes the relevant chunks + user query
-      4) Generate a response from the pipeline
-      5) Add the assistant's response to FAISS as well
+    Chat function with retrieval augmentation.
     """
     pipe = ensure_pipeline()
-
     retriever.add_text(f"User: {user_input}")
-
     top_k = 3
     results = retriever.search(user_input, top_k=top_k)
-
     prompt = build_rag_prompt(user_input, results)
-
     output = pipe(
         prompt,
         temperature=float(temperature),
@@ -331,16 +307,15 @@ def chat_rag(user_input, history, temperature, top_p, min_new_tokens, max_new_to
         assistant_reply = output.strip()
 
     retriever.add_text(f"Assistant: {assistant_reply}")
-
     history.append([user_input, assistant_reply])
     return history, history
 
+# Build the Gradio interface.
 with gr.Blocks() as demo:
-    gr.Markdown("# QLoRA Fine-tuning & RAG-based Chat Demo")
+    gr.Markdown("# QLoRA Fine-tuning & RAG-based Chat Demo using Custom DeepSeekV3 Model")
 
-    finetune_btn = gr.Button("Finetune 4-bit (QLoRA) on Magpie subset (up to 5 min)")
+    finetune_btn = gr.Button("Finetune 4-bit (QLoRA) on ServiceNow-AI/R1-Distill-SFT subset (up to 5 min)")
     status_box = gr.Textbox(label="Finetune Status")
-
     finetune_btn.click(fn=finetune_small_subset, outputs=status_box)
 
     gr.Markdown("## Direct Generation (No Retrieval)")
@@ -349,19 +324,18 @@ with gr.Blocks() as demo:
     top_p = gr.Slider(0.0, 1.0, step=0.05, value=0.9, label="Top-p")
     min_tokens = gr.Slider(1, 2500, value=50, step=10, label="Min New Tokens")
     max_tokens = gr.Slider(1, 2500, value=200, step=50, label="Max New Tokens")
-
-    output_box = gr.Textbox(label="myr1 Output", lines=8)
-    gen_btn = gr.Button("Generate with myr1")
+    output_box = gr.Textbox(label="DeepSeekV3 Output", lines=8)
+    gen_btn = gr.Button("Generate with DeepSeekV3")
     gen_btn.click(
         fn=predict,
         inputs=[prompt_in, temperature, top_p, min_tokens, max_tokens],
         outputs=output_box
     )
 
-    gr.Markdown("## Compare myr1 vs DeepSeek")
+    gr.Markdown("## Compare DeepSeekV3 vs Reference DeepSeek")
     compare_btn = gr.Button("Compare")
-    out_local = gr.Textbox(label="myr1 Output", lines=6)
-    out_deepseek = gr.Textbox(label="DeepSeek Output", lines=6)
+    out_local = gr.Textbox(label="DeepSeekV3 Output", lines=6)
+    out_deepseek = gr.Textbox(label="Reference DeepSeek Output", lines=6)
     compare_btn.click(
         fn=compare_models,
         inputs=[prompt_in, temperature, top_p, min_tokens, max_tokens],
@@ -372,15 +346,13 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             chatbot = gr.Chatbot(label="RAG Chat")
-            chat_state = gr.State([])  
-
+            chat_state = gr.State([])
             user_input = gr.Textbox(
                 show_label=False,
                 placeholder="Ask a question...",
                 lines=2
             )
             send_btn = gr.Button("Send")
-
     user_input.submit(
         fn=chat_rag,
         inputs=[user_input, chat_state, temperature, top_p, min_tokens, max_tokens],
@@ -392,4 +364,4 @@ with gr.Blocks() as demo:
         outputs=[chat_state, chatbot]
     )
 
-demo.launch()
+demo.launch()