Update app.py

app.py

@@ -1,112 +1,25 @@
 import spaces
-import os
-import uuid
-import time
-import json
-import torch
-from datetime import datetime, timedelta
-from threading import Thread
-from pathlib import Path
-
-# Gradio and HuggingFace imports
-import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from datasets import Dataset
-from huggingface_hub import HfApi, login
+import gradio as gr
+from threading import Thread
+import os
+from gradio_modal import Modal
 
-# Model configuration
 checkpoint = "WillHeld/soft-raccoon"
-
-# Set device based on availability
-if torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
-    print("CUDA not available, using CPU instead. This will be much slower.")
-
-# Dataset configuration
-DATASET_NAME = "WillHeld/soft-raccoon-conversations"  # Change to your username
-SAVE_INTERVAL_MINUTES = 5  # Save data every 5 minutes
-last_save_time = datetime.now()
-
-# Initialize model and tokenizer
-print(f"Loading model from {checkpoint}...")
+device = "cuda"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
 
-# Data storage
-conversations = []
-
-# 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, f"No conversations to save. Last attempt: {datetime.now().strftime('%H:%M:%S')}"
-    
-    # Convert conversations to dataset format
-    dataset_dict = {
-        "conversation_id": [],
-        "timestamp": [],
-        "messages": [],
-        "metadata": []
-    }
-    
-    for conv in conversations:
-        dataset_dict["conversation_id"].append(conv["conversation_id"])
-        dataset_dict["timestamp"].append(conv["timestamp"])
-        dataset_dict["messages"].append(json.dumps(conv["messages"]))
-        dataset_dict["metadata"].append(json.dumps(conv["metadata"]))
-    
-    # Create dataset
-    dataset = Dataset.from_dict(dataset_dict)
-    
-    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, status_msg
-
 @spaces.GPU(duration=120)
-def chat_model(message, history, temperature=0.7, top_p=0.9):
-    """Chat function for use with ChatInterface"""
-    conversation_id = getattr(chat_model, "conversation_id", None)
-    if conversation_id is None:
-        conversation_id = str(uuid.uuid4())
-        chat_model.conversation_id = conversation_id
-    
-    # Format chat history for the model
-    formatted_history = []
-    for h in history:
-        formatted_history.append({"role": "user", "content": h["content"] if h["role"] == "user" else ""})
-        if h["role"] == "assistant":
-            formatted_history.append({"role": "assistant", "content": h["content"]})
-    
-    # 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
-    )
+def predict(message, history, temperature, top_p):
+    history.append({"role": "user", "content": message})
+    input_text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
     
-    # Set up streaming
+    # Create a streamer
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Generation parameters
+    # Set up generation parameters
     generation_kwargs = {
         "input_ids": inputs,
         "max_new_tokens": 1024,
@@ -114,210 +27,73 @@ def chat_model(message, history, temperature=0.7, top_p=0.9):
         "top_p": float(top_p),
         "do_sample": True,
         "streamer": streamer,
+        "eos_token_id": 128009,
     }
     
-    # Generate in a separate thread
+    # Run generation in a separate thread
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    # Initialize response
+    # Yield from the streamer as tokens are generated
     partial_text = ""
-    
-    # Yield partial text as it's generated
     for new_text in streamer:
         partial_text += new_text
         yield partial_text
-    
-    # Store conversation data in the global conversations list
-    formatted_history.append({"role": "assistant", "content": partial_text})
-    
-    # Find existing conversation or create new one
-    existing_conv = next((c for c in conversations if c["conversation_id"] == conversation_id), None)
-    
-    # Update or create conversation record
-    current_time = datetime.now().isoformat()
-    if existing_conv:
-        # Update existing conversation
-        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": current_time,
-            "messages": formatted_history,
-            "metadata": {
-                "model": checkpoint,
-                "temperature": temperature,
-                "top_p": top_p,
-                "last_updated": current_time
-            }
-        })
-    
-    # Check if it's time to save based on elapsed time
-    global last_save_time
-    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_dt
-
-def save_dataset_manually():
-    """Manually trigger dataset save and return status"""
-    _, 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 {
-        "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
-    }
-
-def update_stats_event():
-    """Update the stats UI elements"""
-    stats = get_stats()
-    return [
-        stats["conversation_count"],
-        stats["next_save"],
-        stats["last_save"],
-        stats["dataset_name"]
-    ]
 
-# Create a Stanford theme
-theme = gr.themes.Default(
-    primary_hue=gr.themes.utils.colors.red,
-    secondary_hue=gr.themes.utils.colors.gray,
-    neutral_hue=gr.themes.utils.colors.gray,
-    font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui"]
-).set(
-    button_primary_background_fill="#8C1515",
-    button_primary_background_fill_hover="#771212",
-    button_primary_text_color="white",
-    slider_color="#8C1515",
-    block_title_text_color="#8C1515",
-    block_label_text_color="#4D4F53",
-    input_border_color_focus="#8C1515",
-    checkbox_background_color_selected="#8C1515",
-    checkbox_border_color_selected="#8C1515",
-    button_secondary_border_color="#4D4F53",
-    block_title_background_fill="#f5f5f5",
-    block_label_background_fill="#f9f9f9"
-)
+# Function to handle the report submission
+def submit_report(satisfaction, feedback_text):
+    # In a real application, you might save this to a database or file
+    print(f"Report submitted - Satisfaction: {satisfaction}, Feedback: {feedback_text}")
+    return "Thank you for your feedback! Your report has been submitted."
 
-# Custom CSS
-css = """
-.gradio-container {
-    font-family: 'Source Sans Pro', sans-serif !important;
-}
-.footer {
-    color: #4D4F53 !important;
-    font-size: 0.85em !important;
-}
-"""
-
-# Set up the Gradio app with Blocks for more control
-with gr.Blocks(theme=theme, title="Stanford Soft Raccoon Chat", css=css) as demo:
-    # Create a timer for periodic updates
-    timer = gr.Timer(30, update_stats_event)  # Update every 30 seconds
-    
+with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=3):
-            # Create the chatbot component directly
-            chatbot = gr.Chatbot(
-                label="Soft Raccoon Chat",
-                avatar_images=(None, "🌲"),  # Stanford tree emoji
-                height=600,
-                placeholder="<strong>Soft Raccoon AI Assistant</strong><br>Ask me anything!"
-            )
-            
-            # Create sliders for temperature and top_p
-            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"
-                )
-            
-            # Create the ChatInterface
-            chat_interface = gr.ChatInterface(
-                fn=chat_model,
-                chatbot=chatbot,
-                additional_inputs=[temperature, top_p],
-                type="messages",  # This is important for compatibility
-                title="Stanford Soft Raccoon Chat",
-                description="AI assistant powered by the Soft Raccoon language model",
-                examples=[
-                    ["Tell me about Stanford University", 0.7, 0.9],
-                    ["How can I learn about artificial intelligence?", 0.8, 0.95],
-                    ["What's your favorite book?", 0.6, 0.85]
+            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")
                 ],
-                cache_examples=True,
+                type="messages"
             )
         
         with gr.Column(scale=1):
-            with gr.Group():
-                gr.Markdown("### Dataset Controls")
-                save_button = gr.Button("Save conversations now", variant="secondary")
-                status_output = gr.Textbox(label="Save Status", interactive=False)
-                
-                with gr.Row():
-                    convo_count = gr.Number(label="Total Conversations", interactive=False)
-                    next_save = gr.Textbox(label="Next Auto-Save", interactive=False)
-                
-                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_manually,
-        [],
-        [status_output]
-    )
-    
-    refresh_btn.click(
-        update_stats_event,
-        [],
-        [convo_count, next_save, last_save_time_display, dataset_name_display]
-    )
+            report_button = gr.Button("File a Report", variant="primary")
     
-    # Connect the timer to update the stats displays
-    timer.stream(
-        fn=lambda: None,  # No-op function
-        inputs=None,
-        outputs=None,
+    # Create the modal with feedback form components
+    with Modal(visible=False) as feedback_modal:
+        with gr.Column():
+            gr.Markdown("## We value your feedback!")
+            gr.Markdown("Please tell us about your experience with the model.")
+            
+            satisfaction = gr.Radio(
+                ["Very satisfied", "Satisfied", "Neutral", "Unsatisfied", "Very unsatisfied"],
+                label="How satisfied are you with the model's responses?",
+                value="Neutral"
+            )
+            
+            feedback_text = gr.Textbox(
+                lines=5,
+                label="Please provide any additional feedback or describe issues you encountered:",
+                placeholder="Enter your detailed feedback here..."
+            )
+            
+            submit_button = gr.Button("Submit Feedback", variant="primary")
+            response_text = gr.Textbox(label="Status", interactive=False)
+    
+    # Connect the "File a Report" button to show the modal
+    report_button.click(
+        lambda: Modal(visible=True),
+        None,
+        feedback_modal
     )
     
-    # Initial update of stats on page load
-    demo.load(
-        update_stats_event,
-        [],
-        [convo_count, next_save, last_save_time_display, dataset_name_display]
+    # Connect the submit button to the submit_report function
+    submit_button.click(
+        submit_report,
+        inputs=[satisfaction, feedback_text],
+        outputs=response_text
     )
-    
-    # Ensure we save on shutdown
-    import atexit
-    atexit.register(save_to_dataset)
 
-# Launch the app
-if __name__ == "__main__":
-    demo.queue()  # Enable queuing for better performance
-    demo.launch(share=True)
+demo.launch()