Update app.py

app.py

@@ -1,38 +1,43 @@
-import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-import gradio as gr
-from threading import Thread
+import os
+import uuid
+import time
+import json
 from datetime import datetime, timedelta
+from threading import Thread
+
+# Gradio and HuggingFace imports
+import gradio as gr
+from gradio.themes import Base
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from datasets import Dataset
 from huggingface_hub import HfApi, login
-import uuid
-import os
-import time
 
+# Model configuration
 checkpoint = "WillHeld/soft-raccoon"
 device = "cuda"
-tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 
 # Dataset configuration
-DATASET_NAME = "WillHeld/soft-raccoon-conversations"  # Change to your HF username
-PUSH_TO_HUB = True  # Set to False if you just want to save locally first
-
-# Time-based storage settings
-SAVE_INTERVAL_MINUTES = 5  # Save every 5 minutes
+DATASET_NAME = "your-username/soft-raccoon-conversations"  # Change to your username
+SAVE_INTERVAL_MINUTES = 5  # Save data every 5 minutes
 last_save_time = datetime.now()
 
-# Initialize storage for conversations
+# Initialize model and tokenizer
+print(f"Loading model from {checkpoint}...")
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
+
+# Data storage
 conversations = []
 
-# Login to Huggingface Hub (you'll need to set HF_TOKEN env var or use login())
-# Uncomment the below line to login with your token
-login(token=os.environ.get("HF_TOKEN"))
+# Hugging Face authentication
+# Uncomment this line to login with your token
+# login(token=os.environ.get("HF_TOKEN"))
+
 
 def save_to_dataset():
     """Save the current conversations to a HuggingFace dataset"""
     if not conversations:
-        return None
+        return None, f"No conversations to save. Last attempt: {datetime.now().strftime('%H:%M:%S')}"
     
     # Convert conversations to dataset format
     dataset_dict = {
@@ -45,43 +50,55 @@ def save_to_dataset():
     for conv in conversations:
         dataset_dict["conversation_id"].append(conv["conversation_id"])
         dataset_dict["timestamp"].append(conv["timestamp"])
-        dataset_dict["messages"].append(conv["messages"])
-        dataset_dict["metadata"].append(conv["metadata"])
+        dataset_dict["messages"].append(json.dumps(conv["messages"]))
+        dataset_dict["metadata"].append(json.dumps(conv["metadata"]))
     
     # Create dataset
     dataset = Dataset.from_dict(dataset_dict)
     
-    if PUSH_TO_HUB:
-        try:
-            # Push to hub - will create the dataset if it doesn't exist
-            dataset.push_to_hub(DATASET_NAME)
-            print(f"Successfully pushed {len(conversations)} conversations to {DATASET_NAME}")
-        except Exception as e:
-            print(f"Error pushing to hub: {e}")
-            # Save locally as fallback
-            dataset.save_to_disk("local_dataset")
-    else:
-        # Save locally
-        dataset.save_to_disk("local_dataset")
-        print(f"Saved {len(conversations)} conversations locally to 'local_dataset'")
+    try:
+        # Push to hub
+        dataset.push_to_hub(DATASET_NAME)
+        status_msg = f"Successfully saved {len(conversations)} conversations to {DATASET_NAME}"
+        print(status_msg)
+    except Exception as e:
+        # Save locally as fallback
+        local_path = f"local_dataset_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        dataset.save_to_disk(local_path)
+        status_msg = f"Error pushing to hub: {str(e)}. Saved locally to '{local_path}'"
+        print(status_msg)
     
-    return dataset
+    return dataset, status_msg
+
 
-@spaces.GPU(duration=120)
-def predict(message, history, temperature, top_p, conversation_id=None):
-    # Create or retrieve conversation ID for tracking
-    if conversation_id is None:
+def predict(message, chat_history, temperature, top_p, conversation_id=None):
+    """Generate a response using the model and save the conversation"""
+    # Create/retrieve conversation ID for tracking
+    if conversation_id is None or conversation_id == "":
         conversation_id = str(uuid.uuid4())
     
-    # Update history with user message
-    history.append({"role": "user", "content": message})
-    input_text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
+    # Format chat history for the model
+    formatted_history = []
+    for human_msg, ai_msg in chat_history:
+        formatted_history.append({"role": "user", "content": human_msg})
+        if ai_msg:  # Skip None values that might occur during streaming
+            formatted_history.append({"role": "assistant", "content": ai_msg})
+    
+    # Add the current message
+    formatted_history.append({"role": "user", "content": message})
+    
+    # Prepare input for the model
+    input_text = tokenizer.apply_chat_template(
+        formatted_history, 
+        tokenize=False,
+        add_generation_prompt=True
+    )
     inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
     
-    # Create a streamer
+    # Set up streaming
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Set up generation parameters
+    # Generation parameters
     generation_kwargs = {
         "input_ids": inputs,
         "max_new_tokens": 1024,
@@ -91,102 +108,200 @@ def predict(message, history, temperature, top_p, conversation_id=None):
         "streamer": streamer,
     }
     
-    # Run generation in a separate thread
+    # Generate in a separate thread
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    # Yield from the streamer as tokens are generated
+    # Initialize response
     partial_text = ""
+    
+    # Yield partial text as it's generated
     for new_text in streamer:
         partial_text += new_text
-        yield partial_text
-    
-    # After generation completes, update history with assistant response
-    history.append({"role": "assistant", "content": partial_text})
+        yield chat_history + [[message, partial_text]], conversation_id
     
     # Store conversation data
-    # Check if we already have this conversation
     existing_conv = next((c for c in conversations if c["conversation_id"] == conversation_id), None)
     
+    # Update history with final response
+    formatted_history.append({"role": "assistant", "content": partial_text})
+    
+    # Update or create conversation record
+    current_time = datetime.now().isoformat()
     if existing_conv:
         # Update existing conversation
-        existing_conv["messages"] = history
-        existing_conv["metadata"]["last_updated"] = datetime.now().isoformat()
+        existing_conv["messages"] = formatted_history
+        existing_conv["metadata"]["last_updated"] = current_time
+        existing_conv["metadata"]["temperature"] = temperature
+        existing_conv["metadata"]["top_p"] = top_p
     else:
         # Create new conversation record
         conversations.append({
             "conversation_id": conversation_id,
-            "timestamp": datetime.now().isoformat(),
-            "messages": history,
+            "timestamp": current_time,
+            "messages": formatted_history,
             "metadata": {
                 "model": checkpoint,
                 "temperature": temperature,
                 "top_p": top_p,
-                "last_updated": datetime.now().isoformat()
+                "last_updated": current_time
             }
         })
     
     # Check if it's time to save based on elapsed time
     global last_save_time
-    current_time = datetime.now()
-    if current_time - last_save_time > timedelta(minutes=SAVE_INTERVAL_MINUTES):
+    current_time_dt = datetime.now()
+    if current_time_dt - last_save_time > timedelta(minutes=SAVE_INTERVAL_MINUTES):
         save_to_dataset()
-        last_save_time = current_time
+        last_save_time = current_time_dt
+    
+    return chat_history + [[message, partial_text]], conversation_id
+
+
+def save_dataset_manually():
+    """Manually trigger dataset save"""
+    _, status = save_to_dataset()
+    return status
+
+
+def get_stats():
+    """Get current stats about conversations and saving"""
+    mins_until_save = SAVE_INTERVAL_MINUTES - (datetime.now() - last_save_time).seconds // 60
+    if mins_until_save < 0:
+        mins_until_save = 0
     
-    return partial_text
+    return {
+        "conversation_count": len(conversations),
+        "next_save": f"In {mins_until_save} minutes",
+        "last_save": last_save_time.strftime('%H:%M:%S'),
+        "dataset_name": DATASET_NAME
+    }
+
+
+# Create a custom Stanford theme
+class StanfordTheme(gr.Theme):
+    def __init__(self):
+        super().__init__(
+            primary_hue={"name": "cardinal", "c50": "#F9E8E8", "c100": "#F0C9C9", "c200": "#E39B9B", 
+                         "c300": "#D66E6E", "c400": "#C94A4A", "c500": "#B82C2C", "c600": "#8C1515", 
+                         "c700": "#771212", "c800": "#620E0E", "c900": "#4D0A0A", "c950": "#380707"},
+            secondary_hue={"name": "cool_gray", "c50": "#F5F5F6", "c100": "#E6E7E8", "c200": "#CDCED0", 
+                           "c300": "#B3B5B8", "c400": "#9A9CA0", "c500": "#818388", "c600": "#4D4F53", 
+                           "c700": "#424448", "c800": "#36383A", "c900": "#2E2D29", "c950": "#1D1D1B"},
+            neutral_hue="gray",
+            radius_size=gr.themes.sizes.radius_sm,
+            font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui"]
+        )
 
-def save_dataset_button():
-    """Manually save the current dataset"""
-    dataset = save_to_dataset()
-    if dataset:
-        return f"Saved {len(conversations)} conversations to dataset."
-    return "No conversations to save."
+# Use the Stanford theme
+theme = StanfordTheme()
 
-with gr.Blocks() as demo:
-    conversation_id = gr.State(None)
+# Set up the Gradio app
+with gr.Blocks(theme=theme, title="Stanford Soft Raccoon Chat with Dataset Collection") as demo:
+    conversation_id = gr.State("")
     
     with gr.Row():
         with gr.Column(scale=3):
-            chatbot = gr.ChatInterface(
-                predict,
-                additional_inputs=[
-                    gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
-                    gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-P"),
-                    conversation_id
-                ],
-                type="messages"
+            chatbot = gr.Chatbot(
+                label="Soft Raccoon Chat",
+                avatar_images=(None, "🦝"),
+                height=600
             )
+            
+            with gr.Row():
+                msg = gr.Textbox(
+                    placeholder="Send a message...",
+                    show_label=False,
+                    container=False
+                )
+                submit_btn = gr.Button("Send", variant="primary")
+            
+            with gr.Accordion("Generation Parameters", open=False):
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=2.0,
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05,
+                    label="Top-P"
+                )
         
         with gr.Column(scale=1):
             with gr.Group():
                 gr.Markdown("### Dataset Controls")
-                save_button = gr.Button("Save conversations to dataset")
-                save_output = gr.Textbox(label="Save Status")
+                save_button = gr.Button("Save conversations now", variant="secondary")
+                status_output = gr.Textbox(label="Save Status", interactive=False)
                 
-                # Display current conversation count
-                conversation_count = gr.Number(value=lambda: len(conversations), 
-                                              label="Total Conversations", 
-                                              interactive=False)
+                with gr.Row():
+                    convo_count = gr.Number(label="Total Conversations", interactive=False)
+                    next_save = gr.Textbox(label="Next Auto-Save", interactive=False)
                 
-                # Display time until next auto-save
-                next_save_time = gr.Textbox(label="Next Auto-Save", 
-                                           value=lambda: f"In {SAVE_INTERVAL_MINUTES - (datetime.now() - last_save_time).seconds // 60} minutes")
-                refresh_button = gr.Button("Refresh Stats")
+                last_save_time_display = gr.Textbox(label="Last Save Time", interactive=False)
+                dataset_name_display = gr.Textbox(label="Dataset Name", interactive=False)
+                
+                refresh_btn = gr.Button("Refresh Stats")
     
     # Set up event handlers
-    save_button.click(save_dataset_button, outputs=save_output)
+    submit_btn.click(
+        predict,
+        [msg, chatbot, temperature, top_p, conversation_id],
+        [chatbot, conversation_id],
+        api_name="chat"
+    )
+    
+    msg.submit(
+        predict,
+        [msg, chatbot, temperature, top_p, conversation_id],
+        [chatbot, conversation_id],
+        api_name=False
+    )
+    
+    save_button.click(
+        save_dataset_manually,
+        [],
+        [status_output]
+    )
+    
+    def update_stats():
+        stats = get_stats()
+        return [
+            stats["conversation_count"],
+            stats["next_save"],
+            stats["last_save"],
+            stats["dataset_name"]
+        ]
+    
+    refresh_btn.click(
+        update_stats,
+        [],
+        [convo_count, next_save, last_save_time_display, dataset_name_display]
+    )
     
-    def refresh_stats():
-        mins_until_save = SAVE_INTERVAL_MINUTES - (datetime.now() - last_save_time).seconds // 60
-        return len(conversations), f"In {mins_until_save} minutes"
+    # Auto-update stats every 30 seconds
+    gr.on(
+        [demo.load, gr.Timeout(30)], 
+        update_stats, 
+        [], 
+        [convo_count, next_save, last_save_time_display, dataset_name_display]
+    )
     
-    refresh_button.click(refresh_stats, outputs=[conversation_count, next_save_time])
+    # Ensure we save on shutdown using atexit
+    import atexit
+    atexit.register(save_to_dataset)
     
-    # Save on shutdown
-    demo.on_close(save_to_dataset)
+    # Set up a function that will be called when the demo loads
+    def on_startup():
+        return update_stats()
     
-    # Set up periodic UI refresh (every 60 seconds)
-    gr.Timer(60, lambda: None).start()
+    demo.load(on_startup, [], [convo_count, next_save, last_save_time_display, dataset_name_display])
 
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)
+