Update app.py

app.py

@@ -3,13 +3,95 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 import gradio as gr
 from threading import Thread
 import os
+import json
+import uuid
+from datasets import Dataset
+from huggingface_hub import HfApi, login
+import huggingface_hub
+import time
+
 from gradio_modal import Modal
+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/marin-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)
+
+# 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
+
 @spaces.GPU(duration=120)
 def predict(message, history, temperature, top_p):
     history.append({"role": "user", "content": message})
@@ -27,7 +109,6 @@ def predict(message, history, temperature, top_p):
         "top_p": float(top_p),
         "do_sample": True,
         "streamer": streamer,
-        "eos_token_id": 128009,
     }
     
     # Run generation in a separate thread
@@ -40,13 +121,30 @@ def predict(message, history, temperature, top_p):
         partial_text += new_text
         yield partial_text
 
-# 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."
+# Function to handle the research feedback submission
+def submit_research_feedback(conversation_history, satisfaction, feedback_text, hf_token=None):
+    """Save user feedback both locally and to HuggingFace Hub"""
+    # Save locally first
+    feedback_id = save_feedback_locally(conversation_history, satisfaction, feedback_text)
+    
+    # Get token from environment variable
+    env_token = os.environ.get("HF_TOKEN")
+    
+    # Use environment token, ignoring any passed 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
 
+# Create the Gradio interface
 with gr.Blocks() as demo:
+    # Store conversation history
+    conversation_state = gr.State([])
+    
     with gr.Row():
         with gr.Column(scale=3):
             chatbot = gr.ChatInterface(
@@ -58,28 +156,43 @@ with gr.Blocks() as demo:
                 type="messages"
             )
         
+            # Update conversation_state with each new message
+            chatbot.submit_btn.click(
+                lambda history: history,
+                inputs=[chatbot.chat_history],
+                outputs=[conversation_state]
+            )
+        
         with gr.Column(scale=1):
-            report_button = gr.Button("File a Report", variant="primary")
+            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("## We value your feedback!")
-            gr.Markdown("Please tell us about your experience with the model.")
+            gr.Markdown("## Research Preview Feedback")
+            gr.Markdown("Thank you for testing our research model. Your feedback (positive or negative) helps us improve!")
+            
+            # Optional: HF Token for pushing to Hub
+            hf_token_input = gr.Textbox(
+                label="HuggingFace Token (Optional)",
+                placeholder="Enter your HF token to push feedback to dataset",
+                type="password",
+                visible=True  # Set to False in production if using environment variables
+            )
             
             satisfaction = gr.Radio(
                 ["Very satisfied", "Satisfied", "Neutral", "Unsatisfied", "Very unsatisfied"],
-                label="How satisfied are you with the model's responses?",
+                label="How would you rate your experience with this research model?",
                 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..."
+                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 Feedback", variant="primary")
+            submit_button = gr.Button("Submit Research Feedback", variant="primary")
             response_text = gr.Textbox(label="Status", interactive=False)
     
     # Connect the "File a Report" button to show the modal
@@ -89,11 +202,12 @@ with gr.Blocks() as demo:
         feedback_modal
     )
     
-    # Connect the submit button to the submit_report function
+    # Connect the submit button to the submit_research_feedback function
     submit_button.click(
-        submit_report,
-        inputs=[satisfaction, feedback_text],
+        submit_research_feedback,
+        inputs=[conversation_state, satisfaction, feedback_text, hf_token_input],
         outputs=response_text
     )
 
+# Launch the demo
 demo.launch()