Spaces:
Runtime error
Runtime error
File size: 2,697 Bytes
ff1152e |
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 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load the model and tokenizer once to avoid reloading
model_name = "moonshotai/Kimi-K2-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def chat_with_kimi_k2(input_text, chat_history=None):
if chat_history is None:
chat_history = []
# Append user input to chat history
chat_history.append(f"User: {input_text}")
# Prepare prompt
prompt = "\n".join(chat_history) + "\nAI:"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract AI response
response_text = response.split("AI:")[-1].strip()
chat_history.append(f"AI: {response_text}")
return response_text, chat_history
# Main menu interface
def main_menu():
return gr.Interface(
fn=show_model_selection,
title="AI Models Hub",
description="Select an AI model to chat with:",
inputs=[],
outputs=gr.Button("Kimi K2"),
live=True
)
# Show model options
def show_model_selection():
with gr.Row():
kimi_button = gr.Button("Chat with Kimi K2")
return kimi_button
# Chat interface for Kimi K2
def kimi_chat():
chat_history = []
def respond(user_input, history):
response, chat_history_updated = chat_with_kimi_k2(user_input, chat_history=history)
return response, chat_history_updated
with gr.Blocks():
gr.Markdown("### Chat with Kimi K2")
chatbox = gr.Textbox(placeholder="Type your message...", label="You")
response_box = gr.Textbox(label="Kimi K2")
chat_history_state = gr.State([])
back_button = gr.Button("Back to Menu")
def back():
return "Back", None
back_button.click(fn=main_menu, inputs=None, outputs=None)
user_input = chatbox
user_input.submit(respond, [user_input, chat_history_state], [response_box, chat_history_state])
back_button.click(lambda: None, None, None)
return gr.Column([chatbox, response_box, back_button])
# Assemble the main app
def app():
with gr.Blocks() as demo:
menu = gr.Button("Go to Main Menu")
kimichat = gr.Button("Chat with Kimi K2")
def show_kimi():
return kimi_chat()
menu.click(show_kimi)
kimichat.click(show_kimi)
return demo
# Run the app
demo = app()
demo.launch() |