Create app.py

app.py

@@ -0,0 +1,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()