Update app.py

app.py

@@ -1,4 +1,7 @@
-# Add this in your imports section if not already present
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import gradio as gr
+from threading import Thread
 import os
 import json
 import uuid
@@ -6,14 +9,156 @@ from datasets import 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)
+
+# 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
+
+# Create a State to store chat history
+chat_history_state = []
+
+@spaces.GPU(duration=120)
+def predict(message, history, temperature, top_p):
+    global chat_history_state
+    
+    # Update our chat history state
+    history.append({"role": "user", "content": message})
+    chat_history_state = history.copy()
+    
+    input_text = tokenizer.apply_chat_template(history, 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
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        yield partial_text
+    
+    # After generation is complete, update chat history state with the assistant response
+    chat_history_state.append({"role": "assistant", "content": partial_text})
+
+# Function to handle the research feedback submission
+def submit_research_feedback(satisfaction, feedback_text):
+    """Save user feedback both locally and to HuggingFace Hub"""
+    global chat_history_state
+    
+    # Save locally first
+    feedback_id = save_feedback_locally(chat_history_state, 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
+
 # Create the Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=3):
-            # Create a State component to store the conversation history
-            chat_history = gr.State([])
-            
-            # Create the ChatInterface 
             chatbot = gr.ChatInterface(
                 predict,
                 additional_inputs=[
@@ -22,15 +167,6 @@ with gr.Blocks() as demo:
                 ],
                 type="messages"
             )
-            
-            # Create a function to update the chat history state
-            def update_history(message, history):
-                chat_history.value = history
-                return message, history
-            
-            # Intercept chatbot responses to update our history state
-            # This requires modifying your predict function to pass through the history
-            # And connecting it to the update_history function
         
         with gr.Column(scale=1):
             report_button = gr.Button("Share Feedback", variant="primary")
@@ -65,7 +201,10 @@ with gr.Blocks() as demo:
     
     # Connect the submit button to the submit_research_feedback function with the current chat history
     submit_button.click(
-        lambda satisfaction, feedback_text, history: submit_research_feedback(history, satisfaction, feedback_text),
-        inputs=[satisfaction, feedback_text, chat_history],
+        lambda satisfaction, feedback_text: submit_research_feedback(satisfaction, feedback_text),
+        inputs=[satisfaction, feedback_text],
         outputs=response_text
-    )
+    )
+
+# Launch the demo
+demo.launch()