Update app.py

app.py

@@ -2,81 +2,55 @@ import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 import gradio as gr
+from threading import Thread
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL = "AGI-0/Art-v0-3B"
 
+TITLE = """<h2>Link to the model: <a href="https://huggingface.co/AGI-0/Art-v0-3B">click here</a></h2>"""
+
+PLACEHOLDER = """
+<center>
+<p>Hi! How can I help you today?</p>
+</center>
+"""
+
+CSS = """
+.duplicate-button {
+    margin: auto !important;
+    color: white !important;
+    background: black !important;
+    border-radius: 100vh !important;
+}
+h3 {
+    text-align: center;
+}
+"""
+
 class ConversationManager:
     def __init__(self):
-        self.model_messages = []  # Stores raw responses with tags
-        
-    def format_for_display(self, raw_response):
-        """Convert model response to user-friendly markdown.
-        Keeps original response intact for model."""
-        
-        # No response? Return empty
-        if not raw_response:
-            return ""
-            
-        display_response = raw_response
-        
-        # Handle reasoning sections
-        while "<|start_reasoning|>" in display_response and "<|end_reasoning|>" in display_response:
-            start = display_response.find("<|start_reasoning|>")
-            end = display_response.find("<|end_reasoning|>") + len("<|end_reasoning|>")
-            
-            # Extract reasoning content
-            reasoning_block = display_response[start:end]
-            reasoning_content = reasoning_block.replace("<|start_reasoning|>", "").replace("<|end_reasoning|>", "")
-            
-            # Replace with markdown details/summary
-            markdown_block = f"\n<details><summary>View Reasoning</summary>\n\n{reasoning_content}\n\n</details>\n"
-            display_response = display_response[:start] + markdown_block + display_response[end:]
-        
-        # Clean up other tags
-        tags_to_remove = [
-            "<|im_start|>",
-            "<|im_end|>",
-            "<|assistant|>",
-            "<|user|>"
-        ]
-        
-        for tag in tags_to_remove:
-            display_response = display_response.replace(tag, "")
-        
-        # Clean up any extra whitespace
-        display_response = "\n".join(line.strip() for line in display_response.split("\n"))
-        display_response = "\n".join(filter(None, display_response.split("\n")))
-        
-        return display_response.strip()
+        self.user_history = []  # For displaying to user (with markdown)
+        self.model_history = []  # For feeding back to model (with original tags)
     
-    def add_exchange(self, user_message, assistant_response):
-        """Store raw response in model history"""
-        print("\n=== New Exchange ===")
-        print(f"User: {user_message[:100]}{'...' if len(user_message) > 100 else ''}")
-        print(f"Assistant (raw): {assistant_response[:100]}{'...' if len(assistant_response) > 100 else ''}")
-        
-        self.model_messages.append({
-            "role": "user",
-            "content": user_message
-        })
-        self.model_messages.append({
-            "role": "assistant",
-            "content": assistant_response
-        })
-        
-        print(f"Current history length: {len(self.model_messages)} messages")
+    def add_exchange(self, user_message, assistant_response, formatted_response):
+        self.model_history.append((user_message, assistant_response))
+        self.user_history.append((user_message, formatted_response))
+        print(f"\nModel History Exchange:")
+        print(f"User: {user_message}")
+        print(f"Assistant (Original): {assistant_response}")
+        print(f"Assistant (Formatted): {formatted_response}")
     
-    def get_conversation_messages(self):
-        """Get full conversation history for model"""
-        return self.model_messages
+    def get_model_history(self):
+        return self.model_history
+    
+    def get_user_history(self):
+        return self.user_history
 
-# Initialize globals
 conversation_manager = ConversationManager()
-device = "cuda"
+
+device = "cuda"  # for GPU usage or "cpu" for CPU usage
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
@@ -86,6 +60,15 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 end_of_sentence = tokenizer.convert_tokens_to_ids("<|im_end|>")
 
+def format_response(response):
+    """Format the response for user display"""
+    if "<|end_reasoning|>" in response:
+        parts = response.split("<|end_reasoning|>")
+        reasoning = parts[0]
+        rest = parts[1] if len(parts) > 1 else ""
+        return f"<details><summary>Click to see reasoning</summary>\n\n{reasoning}\n\n</details>\n\n{rest}"
+    return response
+
 @spaces.GPU()
 def stream_chat(
     message: str,
@@ -97,34 +80,48 @@ def stream_chat(
     top_k: int = 1,
     penalty: float = 1.1,
 ):
-    print(f"\n=== New Chat Request ===")
-    print(f"Message: {message}")
-    print(f"History length: {len(history)}")
+    print(f'\nNew Chat Request:')
+    print(f'Message: {message}')
+    print(f'History from UI: {history}')
+    print(f'System Prompt: {system_prompt}')
+    print(f'Parameters: temp={temperature}, max_tokens={max_new_tokens}, top_p={top_p}, top_k={top_k}, penalty={penalty}')
     
-    # Build conversation history from model's stored messages
+    # Build conversation from UI history instead of model_history
     conversation = []
-    if system_prompt:
-        conversation.append({"role": "system", "content": system_prompt})
+    for prompt, answer in (history or []):
+        # Extract original response if it's in the details format
+        if "<details>" in answer:
+            # Extract content between <details> tags and after </details>
+            parts = answer.split("</details>")
+            if len(parts) > 1:
+                # Get the content after the </details> tag
+                answer_content = parts[1].strip()
+                # Get the reasoning part
+                reasoning = answer.split("<summary>")[1].split("</summary>")[1].strip()
+                # Reconstruct the original format
+                answer = f"{reasoning}<|end_reasoning|>{answer_content}"
+            else:
+                # If no </details> tag found, use the answer as is
+                answer = answer
+        conversation.extend([
+            {"role": "user", "content": prompt},
+            {"role": "assistant", "content": answer},
+        ])
     
-    # Add all previous messages
-    conversation.extend(conversation_manager.get_conversation_messages())
-    
-    # Add new message
     conversation.append({"role": "user", "content": message})
+    print(f'\nFormatted Conversation for Model:')
+    print(conversation)
     
-    print(f"Sending {len(conversation)} messages to model")
-    
-    # Prepare model input
     input_ids = tokenizer.apply_chat_template(
-        conversation,
-        add_generation_prompt=True,
+        conversation, 
+        add_generation_prompt=True, 
         return_tensors="pt"
     ).to(model.device)
     
     streamer = TextIteratorStreamer(
-        tokenizer,
-        timeout=60.0,
-        skip_prompt=True,
+        tokenizer, 
+        timeout=60.0, 
+        skip_prompt=True, 
         skip_special_tokens=True
     )
     
@@ -140,9 +137,8 @@ def stream_chat(
         streamer=streamer,
     )
     
-    # Storage for building complete response
     buffer = ""
-    model_response = ""
+    original_response = ""
     
     with torch.no_grad():
         thread = Thread(target=model.generate, kwargs=generate_kwargs)
@@ -150,56 +146,86 @@ def stream_chat(
         
         for new_text in streamer:
             buffer += new_text
-            model_response += new_text
+            original_response += new_text
             
-            # Convert current buffer for display
-            display_text = conversation_manager.format_for_display(buffer)
+            formatted_buffer = format_response(buffer)
             
-            if not thread.is_alive():
-                print("Generation complete")
-                # Store final response in model history
-                conversation_manager.add_exchange(message, model_response)
+            if thread.is_alive() is False:
+                print(f'\nGeneration Complete:')
+                print(f'Original Response: {original_response}')
+                print(f'Formatted Response: {formatted_buffer}')
+                
+                conversation_manager.add_exchange(
+                    message,
+                    original_response,  # Original for model
+                    formatted_buffer    # Formatted for user
+                )
             
-            yield display_text
+            yield formatted_buffer
 
-# Set up Gradio interface
-CSS = """
-.duplicate-button {
-    margin: auto !important;
-    color: white !important;
-    background: black !important;
-    border-radius: 100vh !important;
-}
-h3 { text-align: center; }
-"""
-
-chatbot = gr.Chatbot(
-    height=600,
-    placeholder="""
-    <center>
-    <p>Hi! How can I help you today?</p>
-    </center>
-    """
-)
+chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)
 
 with gr.Blocks(css=CSS, theme="soft") as demo:
-    gr.HTML("""<h2>Link to the model: <a href="https://huggingface.co/AGI-0/Art-v0-3B">click here</a></h2>""")
+    gr.HTML(TITLE)
     gr.DuplicateButton(
-        value="Duplicate Space for private use",
+        value="Duplicate Space for private use", 
         elem_classes="duplicate-button"
     )
     gr.ChatInterface(
         fn=stream_chat,
         chatbot=chatbot,
         fill_height=True,
-        additional_inputs_accordion=gr.Accordion("⚙️ Parameters", open=False, render=False),
+        additional_inputs_accordion=gr.Accordion(
+            label="⚙️ Parameters",
+            open=False,
+            render=False
+        ),
         additional_inputs=[
-            gr.Textbox(value="", label="System Prompt", render=False),
-            gr.Slider(minimum=0, maximum=1, step=0.1, value=0.2, label="Temperature", render=False),
-            gr.Slider(minimum=128, maximum=8192, step=1, value=4096, label="Max new tokens", render=False),
-            gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="top_p", render=False),
-            gr.Slider(minimum=1, maximum=50, step=1, value=1, label="top_k", render=False),
-            gr.Slider(minimum=0.0, maximum=2.0, step=0.1, value=1.1, label="Repetition penalty", render=False),
+            gr.Textbox(
+                value="",
+                label="System Prompt",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0,
+                maximum=1,
+                step=0.1,
+                value=0.2,
+                label="Temperature",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=128,
+                maximum=8192,
+                step=1,
+                value=4096,
+                label="Max new tokens",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                step=0.1,
+                value=1.0,
+                label="top_p",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=1,
+                maximum=50,
+                step=1,
+                value=1,
+                label="top_k",
+                render=False,
+            ),
+            gr.Slider(
+                minimum=0.0,
+                maximum=2.0,
+                step=0.1,
+                value=1.1,
+                label="Repetition penalty",
+                render=False,
+            ),
         ],
         examples=[
             ["Help me study vocabulary: write a sentence for me to fill in the blank, and I'll try to pick the correct option."],