app.py

import gradio as gr
from huggingface_hub import InferenceClient, HfApi
from datetime import datetime
import uuid
import os
import json

# ---- Configuration ----
MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
DATASET_REPO = "frimelle/companion-chat-logs"
HF_TOKEN = os.environ.get("HF_TOKEN")

# ---- Load system prompt ----
with open("system_prompt.txt", "r") as f:
    SYSTEM_PROMPT = f.read()

client = InferenceClient(MODEL_NAME)
api = HfApi()

# ---- Fixed generation parameters ----
MAX_TOKENS = 512
TEMPERATURE = 0.7
TOP_P = 0.95

# ---- Logging function ----
def log_to_dataset(user_message, assistant_message):
    row = {
        "timestamp": datetime.now().isoformat(),
        "session_id": str(uuid.uuid4()),
        "user": user_message,
        "assistant": assistant_message,
        "system_prompt": SYSTEM_PROMPT,
    }

    # Save as a JSONL file and upload to the dataset repo
    filename = f"log_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}.jsonl"
    with open(filename, "w", encoding="utf-8") as f:
        f.write(json.dumps(row) + "\n")

    api.upload_file(
        path_or_fileobj=filename,
        path_in_repo=filename,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN
    )

    os.remove(filename)

# ---- Chatbot function ----
def respond(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    
    messages.append({"role": "user", "content": message})
    response = ""

    for chunk in client.chat_completion(
        messages,
        max_tokens=MAX_TOKENS,
        stream=True,
        temperature=TEMPERATURE,
        top_p=TOP_P,
    ):
        token = chunk.choices[0].delta.content
        if token:
            response += token
            yield response

    # Log after full response is generated
    log_to_dataset(message, response)

# ---- Gradio Interface ----
demo = gr.ChatInterface(
    respond,
    title="Zephyr Chatbot",
)

if __name__ == "__main__":
    demo.launch()