sahabat tai

app.py

@@ -1,4 +1,4 @@
-import streamlit as st
+import gradio as gr
 import torch
 import transformers
 
@@ -17,40 +17,50 @@ terminators = [
     pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
 ]
 
-# Streamlit App Configuration
-st.set_page_config(page_title="Chatbot", page_icon="🤗")
-st.title("Gemma2 Chatbot")
-st.markdown("A chatbot that understands Javanese and Sundanese using `GoToCompany/gemma2` model.")
+# Chatbot Functionality
+def chatbot(messages):
+    """
+    Handles user interactions and returns the model's response.
+    Args:
+        messages (list): List of messages with roles ('user' or 'assistant') and content.
 
-# Initialize session state
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
-
-# User input
-user_input = st.chat_input("Type your message here...")
-
-# Generate response
-if user_input:
-    # Add user message to chat history
-    st.session_state.chat_history.append({"role": "user", "content": user_input})
-
-    # Prepare conversation context
-    conversation = [
-        {"role": msg["role"], "content": msg["content"]} for msg in st.session_state.chat_history
-    ]
-
-    # Generate response using the pipeline
+    Returns:
+        list: Updated conversation with the assistant's response.
+    """
+    # Prepare the conversation for the model
     outputs = pipeline(
-        conversation,
+        messages,
         max_new_tokens=256,
         eos_token_id=terminators,
     )
-
     # Extract and format the assistant's response
-    assistant_response = outputs[0]["generated_text"][-1] if outputs else "Sorry, I couldn't generate a response."
-    st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+    assistant_response = outputs[0]["generated_text"] if outputs else "I'm sorry, I couldn't generate a response."
+    messages.append({"role": "assistant", "content": assistant_response})
+    return messages
+
+# Gradio Chat Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤗 Gemma2 Chatbot")
+    gr.Markdown("A chatbot that understands Javanese and Sundanese, powered by `GoToCompany/gemma2`.")
+
+    chat_history = gr.Chatbot(label="Gemma2 Chatbot")
+    user_input = gr.Textbox(label="Your Message", placeholder="Type your message here...")
+    send_button = gr.Button("Send")
+
+    def respond(chat_history, user_message):
+        # Add user message to chat history
+        chat_history.append(("user", user_message))
+
+        # Generate assistant's response
+        conversation = [{"role": role, "content": content} for role, content in chat_history]
+        response = chatbot(conversation)
+
+        # Add assistant's response to chat history
+        assistant_message = response[-1]["content"]
+        chat_history.append(("assistant", assistant_message))
+        return chat_history, ""
+
+    send_button.click(respond, inputs=[chat_history, user_input], outputs=[chat_history, user_input])
 
-# Display the chat history
-for message in st.session_state.chat_history:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+# Launch the app
+demo.launch()