Spaces:
Runtime error
Runtime error
File size: 3,351 Bytes
c61047e 7180099 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import os
import json
import gradio as gr
from datetime import datetime
from huggingface_hub import InferenceClient
# Constants
DEFAULT_MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.1"
HF_TOKEN = os.environ.get("HF_TOKEN")
LOG_FILE = "chat_log.json"
def chat_with_model(system_prompt, user_message, max_tokens=500, model_id=""):
"""
Generate a chat completion using the specified or default Mistral model.
Args:
system_prompt (str): The system prompt to set the context.
user_message (str): The user's input message.
max_tokens (int): Maximum number of tokens to generate.
model_id (str): The model ID to use. If empty, uses the default model.
Returns:
str: The model's response.
"""
model_id = model_id.strip() if model_id else DEFAULT_MODEL_ID
client = InferenceClient(model_id, token=HF_TOKEN)
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message}
]
response = ""
try:
for message in client.chat_completion(
messages=messages,
max_tokens=max_tokens,
stream=True
):
response += message.choices[0].delta.content or ""
except Exception as e:
response = f"Error: {str(e)}"
# Log the input and output
log_chat(system_prompt, user_message, response, max_tokens, model_id)
return response
def log_chat(system_prompt, user_message, response, max_tokens, model_id):
"""
Log the chat details to a JSON file.
"""
log_entry = {
"timestamp": datetime.now().isoformat(),
"system_prompt": system_prompt,
"user_message": user_message,
"response": response,
"max_tokens": max_tokens,
"model_id": model_id
}
try:
# Read existing log file if it exists
if os.path.exists(LOG_FILE):
with open(LOG_FILE, 'r') as f:
log_data = json.load(f)
else:
log_data = []
# Append new entry
log_data.append(log_entry)
# Write updated log back to file
with open(LOG_FILE, 'w') as f:
json.dump(log_data, f, indent=2)
except Exception as e:
print(f"Error logging chat: {str(e)}")
def create_gradio_interface():
"""Create and configure the Gradio interface."""
return gr.Interface(
fn=chat_with_model,
inputs=[
gr.Textbox(label="System Prompt", placeholder="Enter the system prompt here..."),
gr.Textbox(label="User Message", placeholder="Ask a question..."),
gr.Slider(minimum=50, maximum=1000, value=500, step=50, label="Max Tokens"),
gr.Textbox(label="Model ID", placeholder=f"Enter model ID (default: {DEFAULT_MODEL_ID})")
],
outputs=gr.Textbox(label="Response"),
title="Mistral Chatbot",
description="Chat with Mistral model using your own system prompts and choose your model.",
examples=[
["You are a helpful AI assistant.", "What is the capital of France?", 500, ""],
["You are an expert in Python programming.", "Explain list comprehensions.", 500, ""]
]
)
if __name__ == "__main__":
iface = create_gradio_interface()
iface.launch(show_api=True, share=False) |