Update app.py

app.py

@@ -2,175 +2,82 @@ import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
-import re
+import gradio as gr
 
 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.user_history = []  # For displaying to user (markdown)
-        self.model_history = []  # For feeding back to model (special tags)
-        self.debug_log = []
+        self.model_messages = []  # Stores raw responses with tags
         
-    def log(self, message):
-        """Add timestamped log entry"""
-        timestamp = time.strftime('%Y-%m-%d %H:%M:%S')
-        log_entry = f"[{timestamp}] {message}"
-        print(log_entry)
-        self.debug_log.append(log_entry)
-    
-    def convert_to_markdown(self, model_text):
-        """Convert from model format (with special tags) to markdown"""
-        self.log(f"Converting to markdown - Input length: {len(model_text)}")
-        self.log(f"Input text: {model_text[:200]}..." if len(model_text) > 200 else f"Input text: {model_text}")
+    def format_for_display(self, raw_response):
+        """Convert model response to user-friendly markdown.
+        Keeps original response intact for model."""
         
-        markdown_text = model_text
+        # No response? Return empty
+        if not raw_response:
+            return ""
+            
+        display_response = raw_response
         
-        # Convert special tags to markdown
-        tag_conversions = [
-            # Reasoning blocks
-            ("<|start_reasoning|>", "<details><summary>Reasoning</summary>\n\n"),
-            ("<|end_reasoning|>", "\n\n</details>\n\n"),
+        # 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|>", "")
             
-            # Other special tags (add more as needed)
-            ("<|im_start|>", ""),
-            ("<|im_end|>", ""),
-            ("<|assistant|>", ""),
-            ("<|user|>", ""),
+            # 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 old, new in tag_conversions:
-            if old in markdown_text:
-                self.log(f"Converting tag: {old} -> {new}")
-                markdown_text = markdown_text.replace(old, new)
-        
-        # Clean up any remaining special tags using regex
-        markdown_text = re.sub(r'<\|[^>]+\|>', '', markdown_text)
-        
-        # Fix common markdown formatting issues
-        markdown_text = re.sub(r'\n{3,}', '\n\n', markdown_text)  # Remove excess newlines
-        markdown_text = markdown_text.strip()
-        
-        self.log(f"Markdown conversion complete - Output length: {len(markdown_text)}")
-        self.log(f"Output text: {markdown_text[:200]}..." if len(markdown_text) > 200 else f"Output text: {markdown_text}")
-        
-        return markdown_text
-    
-    def convert_to_model_format(self, markdown_text):
-        """Convert from markdown to model format (with special tags)"""
-        self.log(f"Converting to model format - Input length: {len(markdown_text)}")
-        self.log(f"Input text: {markdown_text[:200]}..." if len(markdown_text) > 200 else f"Input text: {markdown_text}")
-        
-        model_text = markdown_text
-        
-        # Convert markdown to special tags
-        if "<details>" in markdown_text and "</details>" in markdown_text:
-            try:
-                # Extract content between details tags
-                pattern = r'<details><summary>.*?</summary>\s*(.*?)\s*</details>'
-                matches = re.findall(pattern, markdown_text, re.DOTALL)
-                
-                for match in matches:
-                    original = f"<details><summary>Reasoning</summary>\n\n{match}\n\n</details>"
-                    replacement = f"<|start_reasoning|>{match}<|end_reasoning|>"
-                    model_text = model_text.replace(original, replacement)
-                    self.log(f"Converted details block to reasoning tags")
-            except Exception as e:
-                self.log(f"Warning: Failed to convert details block: {str(e)}")
+        for tag in tags_to_remove:
+            display_response = display_response.replace(tag, "")
         
-        # Clean up formatting
-        model_text = re.sub(r'\n{3,}', '\n\n', model_text)  # Remove excess newlines
-        model_text = model_text.strip()
+        # 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")))
         
-        self.log(f"Model format conversion complete - Output length: {len(model_text)}")
-        self.log(f"Output text: {model_text[:200]}..." if len(model_text) > 200 else f"Output text: {model_text}")
-        
-        return model_text
+        return display_response.strip()
     
     def add_exchange(self, user_message, assistant_response):
-        """Add a new exchange to both histories"""
-        self.log(f"\n=== Adding New Exchange ===")
-        self.log(f"User Message: {user_message[:100]}..." if len(user_message) > 100 else f"User Message: {user_message}")
-        self.log(f"Assistant Response: {assistant_response[:100]}..." if len(assistant_response) > 100 else f"Assistant Response: {assistant_response}")
-        
-        # Convert assistant response to markdown for user display
-        markdown_response = self.convert_to_markdown(assistant_response)
-        
-        # Store both versions
-        self.model_history.append((user_message, assistant_response))
-        self.user_history.append((user_message, markdown_response))
-        
-        self.log(f"Current History State:")
-        self.log(f"- Model History: {len(self.model_history)} exchanges")
-        self.log(f"- User History: {len(self.user_history)} exchanges")
-    
-    def sync_with_ui_history(self, ui_history):
-        """Sync our histories with the UI history"""
-        self.log(f"\n=== Syncing with UI History ===")
-        self.log(f"UI History Length: {len(ui_history)}")
-        
-        # Clear current histories
-        self.model_history = []
-        self.user_history = []
-        
-        # Rebuild histories from UI
-        for user_msg, markdown_response in ui_history:
-            model_response = self.convert_to_model_format(markdown_response)
-            self.model_history.append((user_msg, model_response))
-            self.user_history.append((user_msg, markdown_response))
-        
-        self.log(f"Sync Complete:")
-        self.log(f"- Model History: {len(self.model_history)} exchanges")
-        self.log(f"- User History: {len(self.user_history)} exchanges")
-        
-        # Verify sync integrity
-        if len(self.model_history) != len(self.user_history) or len(self.model_history) != len(ui_history):
-            self.log("WARNING: History length mismatch after sync!")
-    
-    def get_model_history(self):
-        """Get history in model format"""
-        self.log(f"\nReturning Model History ({len(self.model_history)} exchanges)")
-        return self.model_history
-    
-    def get_user_history(self):
-        """Get history in markdown format"""
-        self.log(f"\nReturning User History ({len(self.user_history)} exchanges)")
-        return self.user_history
-    
-    def get_debug_log(self):
-        """Get the full debug log"""
-        return "\n".join(self.debug_log)
+        """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 get_conversation_messages(self):
+        """Get full conversation history for model"""
+        return self.model_messages
 
-# Initialize global conversation manager
+# Initialize globals
 conversation_manager = ConversationManager()
+device = "cuda"
 
-device = "cuda"  # for GPU usage or "cpu" for CPU usage
-
-# Initialize model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
@@ -190,33 +97,24 @@ def stream_chat(
     top_k: int = 1,
     penalty: float = 1.1,
 ):
-    conversation_manager.log(f'\n=== New Chat Request ===')
-    conversation_manager.log(f'Message: {message}')
-    conversation_manager.log(f'History Length: {len(history)}')
-    conversation_manager.log(f'System Prompt: {system_prompt}')
-    conversation_manager.log(f'Parameters: temp={temperature}, max_tokens={max_new_tokens}, top_p={top_p}, top_k={top_k}, penalty={penalty}')
+    print(f"\n=== New Chat Request ===")
+    print(f"Message: {message}")
+    print(f"History length: {len(history)}")
     
-    # Sync with UI history
-    conversation_manager.sync_with_ui_history(history)
-    
-    # Get model-formatted history
-    model_history = conversation_manager.get_model_history()
-    
-    # Build conversation for model
+    # Build conversation history from model's stored messages
     conversation = []
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
     
-    for prompt, answer in model_history:
-        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})
     
-    conversation_manager.log(f'Built conversation with {len(conversation)} messages')
+    print(f"Sending {len(conversation)} messages to model")
     
+    # Prepare model input
     input_ids = tokenizer.apply_chat_template(
         conversation,
         add_generation_prompt=True,
@@ -242,8 +140,9 @@ def stream_chat(
         streamer=streamer,
     )
     
+    # Storage for building complete response
     buffer = ""
-    original_response = ""
+    model_response = ""
     
     with torch.no_grad():
         thread = Thread(target=model.generate, kwargs=generate_kwargs)
@@ -251,27 +150,40 @@ def stream_chat(
         
         for new_text in streamer:
             buffer += new_text
-            original_response += new_text
+            model_response += new_text
             
-            # Convert buffer to markdown for display
-            formatted_buffer = conversation_manager.convert_to_markdown(buffer)
+            # Convert current buffer for display
+            display_text = conversation_manager.format_for_display(buffer)
             
-            if thread.is_alive() is False:
-                conversation_manager.log(f'Generation Complete:')
-                conversation_manager.log(f'Final Response Length: {len(original_response)}')
-                
-                conversation_manager.add_exchange(
-                    message,
-                    original_response  # Original for model
-                )
+            if not thread.is_alive():
+                print("Generation complete")
+                # Store final response in model history
+                conversation_manager.add_exchange(message, model_response)
             
-            yield formatted_buffer
+            yield display_text
+
+# Set up Gradio interface
+CSS = """
+.duplicate-button {
+    margin: auto !important;
+    color: white !important;
+    background: black !important;
+    border-radius: 100vh !important;
+}
+h3 { text-align: center; }
+"""
 
-# Initialize Gradio interface
-chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)
+chatbot = gr.Chatbot(
+    height=600,
+    placeholder="""
+    <center>
+    <p>Hi! How can I help you today?</p>
+    </center>
+    """
+)
 
 with gr.Blocks(css=CSS, theme="soft") as demo:
-    gr.HTML(TITLE)
+    gr.HTML("""<h2>Link to the model: <a href="https://huggingface.co/AGI-0/Art-v0-3B">click here</a></h2>""")
     gr.DuplicateButton(
         value="Duplicate Space for private use",
         elem_classes="duplicate-button"
@@ -280,57 +192,14 @@ with gr.Blocks(css=CSS, theme="soft") as demo:
         fn=stream_chat,
         chatbot=chatbot,
         fill_height=True,
-        additional_inputs_accordion=gr.Accordion(
-            label="⚙️ Parameters",
-            open=False,
-            render=False
-        ),
+        additional_inputs_accordion=gr.Accordion("⚙️ 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."],