app.py

import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import gradio as gr
from threading import Thread
import os
import json
import uuid
from datasets import Dataset, load_dataset
from huggingface_hub import HfApi, login
import time

# Install required packages if not present
from gradio_modal import Modal
import huggingface_hub
import datasets

# Model setup
checkpoint = "WillHeld/soft-raccoon"
device = "cuda"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)

# Constants for dataset
DATASET_REPO = "WillHeld/model-feedback"  # Replace with your username
DATASET_PATH = "./feedback_data"  # Local path to store feedback
DATASET_FILENAME = "feedback.jsonl"  # Filename for feedback data

# Ensure feedback directory exists
os.makedirs(DATASET_PATH, exist_ok=True)

# Sync existing dataset from Hub if available
def sync_dataset_from_hub():
    """Download existing dataset from Hub and merge with local data"""
    try:
        # Try to get token from environment variable
        hf_token = os.environ.get("HF_TOKEN")
        if hf_token:
            login(token=hf_token)
        
        # Check if the dataset exists on Hub
        api = HfApi()
        try:
            dataset_info = api.dataset_info(DATASET_REPO)
            # Dataset exists, download it
            print(f"Syncing existing dataset from {DATASET_REPO}")
            remote_dataset = load_dataset(DATASET_REPO)
            
            # Convert to list of dictionaries
            remote_data = [item for item in remote_dataset['train']]
            
            # Check if local file exists
            local_file = os.path.join(DATASET_PATH, DATASET_FILENAME)
            local_data = []
            
            if os.path.exists(local_file):
                # Read local data
                with open(local_file, 'r') as f:
                    for line in f:
                        try:
                            local_data.append(json.loads(line))
                        except json.JSONDecodeError:
                            continue
            
            # Merge data (using IDs to avoid duplicates)
            all_items = {}
            for item in remote_data + local_data:
                all_items[item['id']] = item
            
            # Write back merged data
            with open(local_file, 'w') as f:
                for item in all_items.values():
                    f.write(json.dumps(item) + '\n')
            
            print(f"Synced {len(all_items)} feedback items")
            return True
        
        except Exception as e:
            print(f"Dataset {DATASET_REPO} does not exist yet or could not be accessed: {e}")
            return False
            
    except Exception as e:
        print(f"Error syncing dataset: {e}")
        return False

# Call sync on startup
sync_dataset_from_hub()

# Feedback storage functions
def save_feedback_locally(conversation, satisfaction, feedback_text):
    """Save feedback to a local JSONL file"""
    # Create a unique ID for this feedback entry
    feedback_id = str(uuid.uuid4())
    
    # Create a timestamp
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    
    # Prepare the feedback data
    feedback_data = {
        "id": feedback_id,
        "timestamp": timestamp,
        "conversation": conversation,
        "satisfaction": satisfaction,
        "feedback": feedback_text
    }
    
    # Save to local file
    feedback_file = os.path.join(DATASET_PATH, DATASET_FILENAME)
    with open(feedback_file, "a") as f:
        f.write(json.dumps(feedback_data) + "\n")
    
    return feedback_id

def push_feedback_to_hub(hf_token=None):
    """Push the local feedback data to HuggingFace as a dataset"""
    # Check if we have a token
    if hf_token is None:
        # Try to get token from environment variable
        hf_token = os.environ.get("HF_TOKEN")
        if hf_token is None:
            print("No HuggingFace token provided. Cannot push to Hub.")
            return False
    
    try:
        # Login to HuggingFace
        login(token=hf_token)
        
        # Check if we have data to push
        feedback_file = os.path.join(DATASET_PATH, DATASET_FILENAME)
        if not os.path.exists(feedback_file):
            print("No feedback data to push.")
            return False
        
        # Load data from the JSONL file
        with open(feedback_file, "r") as f:
            feedback_data = [json.loads(line) for line in f]
        
        # Create a dataset from the feedback data
        dataset = Dataset.from_list(feedback_data)
        
        # Push to Hub
        dataset.push_to_hub(
            DATASET_REPO,
            private=True  # Set to False if you want the dataset to be public
        )
        
        print(f"Feedback data pushed to {DATASET_REPO} successfully.")
        return True
    
    except Exception as e:
        print(f"Error pushing feedback data to Hub: {e}")
        return False

# Function to handle the research feedback submission
def submit_research_feedback(conv_history, satisfaction, feedback_text):
    """Save user feedback both locally and to HuggingFace Hub"""
    # Print debug information
    print(f"Saving feedback with conversation history containing {len(conv_history)} messages")
    if conv_history and len(conv_history) > 0:
        print(f"First message: {conv_history[0]['role']}: {conv_history[0]['content'][:30]}...")
        print(f"Last message: {conv_history[-1]['role']}: {conv_history[-1]['content'][:30]}...")
    
    # Save locally first
    feedback_id = save_feedback_locally(conv_history, satisfaction, feedback_text)
    
    # Get token from environment variable
    env_token = os.environ.get("HF_TOKEN")
    
    # Use environment token
    push_success = push_feedback_to_hub(env_token)
    
    if push_success:
        status_msg = "Thank you for your valuable feedback! Your insights have been saved to the dataset."
    else:
        status_msg = "Thank you for your feedback! It has been saved locally, but couldn't be pushed to the dataset. Please check server logs."
    
    return status_msg

# Initial state - set up at app start
def initialize_state():
    """Initialize the conversation state - this could load previous sessions or start fresh"""
    return []  # Start with empty conversation history

# Create the Gradio blocks interface
with gr.Blocks() as demo:
    # Create state to store full conversation history with proper initialization
    conv_state = gr.State(initialize_state)
    
    with gr.Row():
        with gr.Column(scale=3):
            # Create a custom predict function that updates our state
            def enhanced_predict(message, history, temperature, top_p, state):
                # Initialize state if needed
                if state is None:
                    state = []
                    print("Initializing empty state")
                
                # Copy history to state if state is empty but history exists
                if len(state) == 0 and len(history) > 0:
                    state = history.copy()
                    print(f"Copied {len(history)} messages from history to state")
                
                # Add user message to state
                state.append({"role": "user", "content": message})
                
                # Process with the model (this doesn't modify the original history)
                input_text = tokenizer.apply_chat_template(state, tokenize=False, add_generation_prompt=True)
                inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
                
                # Create a streamer
                streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
                
                # Set up generation parameters
                generation_kwargs = {
                    "input_ids": inputs,
                    "max_new_tokens": 1024,
                    "temperature": float(temperature),
                    "top_p": float(top_p),
                    "do_sample": True,
                    "streamer": streamer,
                    "eos_token_id": 128009,
                }
                
                # Run generation in a separate thread
                thread = Thread(target=model.generate, kwargs=generation_kwargs)
                thread.start()
                
                # Yield from the streamer as tokens are generated
                response = ""
                for new_text in streamer:
                    response += new_text
                    # For each partial response, yield the text only
                    # We'll update the state after generation is complete
                    yield response
                
                # After generation completes, update our state with the final response
                state.append({"role": "assistant", "content": response})
                
                # Return the updated state
                return state
            
            # Create a wrapper that connects to ChatInterface but also updates our state
            def chat_with_state(message, history, temperature, top_p):
                # This function is what interfaces with the ChatInterface
                nonlocal conv_state
                
                # Access the current state
                current_state = conv_state.value if conv_state.value else []
                
                # Call the main function that generates responses and updates state
                # This is a generator function, so we need to iterate through its outputs
                response_gen = enhanced_predict(message, history, temperature, top_p, current_state)
                
                # For each response, yield it and also update our state at the end
                last_response = None
                for response in response_gen:
                    last_response = response
                    yield response
                
                # After generation is complete, update our state
                if last_response is not None:
                    # Create a full copy of the history plus the new exchange
                    updated_state = []
                    # Add all previous history
                    for msg in history:
                        updated_state.append(msg.copy())
                    # Add new exchange
                    updated_state.append({"role": "user", "content": message})
                    updated_state.append({"role": "assistant", "content": last_response})
                    
                    # Store in our state
                    conv_state.value = updated_state
                    
                    # Debug
                    print(f"Updated conversation state with {len(updated_state)} messages")
                    if updated_state:
                        last_msg = updated_state[-1]
                        print(f"Last message: {last_msg['role']}: {last_msg['content'][:30]}...")
            
            # Create ChatInterface
            chatbot = gr.ChatInterface(
                chat_with_state,
                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")
                ],
                type="messages"
            )
        
        with gr.Column(scale=1):
            report_button = gr.Button("Share Feedback", variant="primary")
    
    # Create the modal with feedback form components
    with Modal(visible=False) as feedback_modal:
        with gr.Column():
            gr.Markdown("## Research Preview Feedback")
            gr.Markdown("Thank you for testing our research model. Your feedback (positive or negative) helps us improve!")
            
            satisfaction = gr.Radio(
                ["Very satisfied", "Satisfied", "Neutral", "Unsatisfied", "Very unsatisfied"],
                label="How would you rate your experience with this research model?",
                value="Neutral"
            )
            
            feedback_text = gr.Textbox(
                lines=5,
                label="Share your observations (strengths, weaknesses, suggestions):",
                placeholder="We welcome both positive feedback and constructive criticism to help improve this research prototype..."
            )
            
            submit_button = gr.Button("Submit Research Feedback", variant="primary")
            response_text = gr.Textbox(label="Status", interactive=False)
    
    # Connect the "Share Feedback" button to show the modal
    report_button.click(
        lambda: Modal(visible=True),
        None,
        feedback_modal
    )
    
    # Connect the submit button to the submit_research_feedback function
    submit_button.click(
        submit_research_feedback,
        inputs=[conv_state, satisfaction, feedback_text],
        outputs=response_text
    )

# Launch the demo
demo.launch()