Spaces:
Sleeping
Sleeping
import streamlit as st | |
import os | |
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel | |
from huggingface_hub import login | |
import torch | |
class VietnameseChatbot: | |
def __init__(self, model_name="tamgrnguyen/Gemma-2-2b-it-Vietnamese-Aesthetic"): | |
""" | |
Initialize the Vietnamese chatbot with a pre-trained model | |
""" | |
self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
self.model = AutoModelForCausalLM.from_pretrained(model_name) | |
# Use GPU if available | |
self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
self.model.to(self.device) | |
def generate_response(self, conversation_history, max_length=200): | |
""" | |
Generate a response based on conversation history | |
""" | |
# Combine conversation history into a single prompt | |
prompt = "\n".join(conversation_history) | |
# Tokenize input | |
inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device) | |
# Generate response | |
outputs = self.model.generate( | |
inputs.input_ids, | |
max_length=max_length, | |
num_return_sequences=1, | |
no_repeat_ngram_size=2, | |
temperature=0.7, | |
top_k=50, | |
top_p=0.95 | |
) | |
# Decode the generated response | |
response = self.tokenizer.decode(outputs[0], skip_special_tokens=True) | |
# Extract only the new part of the response | |
response = response[len(prompt):].strip() | |
return response | |
def main(): | |
st.set_page_config(page_title="IOGPT", layout="wide") | |
st.title("Chat với IOGPT") | |
st.markdown(""" | |
### Trò chuyện với IOGPT | |
""") | |
# Initialize chatbot | |
if 'chatbot' not in st.session_state: | |
try: | |
st.session_state.chatbot = VietnameseChatbot() | |
except Exception as e: | |
st.error(f"Lỗi khởi tạo mô hình: {e}") | |
return | |
# Initialize conversation history | |
if 'conversation' not in st.session_state: | |
st.session_state.conversation = [] | |
# Chat interface | |
with st.form(key='chat_form'): | |
user_input = st.text_input("Nhập tin nhắn của bạn:", placeholder="Viết tin nhắn...") | |
send_button = st.form_submit_button("Gửi") | |
# Process user input | |
if send_button and user_input: | |
# Add user message to conversation | |
st.session_state.conversation.append(f"User: {user_input}") | |
try: | |
# Generate AI response | |
with st.spinner('Đang suy nghĩ...'): | |
ai_response = st.session_state.chatbot.generate_response( | |
st.session_state.conversation | |
) | |
# Add AI response to conversation | |
st.session_state.conversation.append(f"AI: {ai_response}") | |
except Exception as e: | |
st.error(f"Lỗi trong quá trình trò chuyện: {e}") | |
# Display conversation history | |
st.subheader("Lịch sử trò chuyện") | |
for msg in st.session_state.conversation: | |
if msg.startswith("User:"): | |
st.markdown(f"**{msg}**") | |
else: | |
st.markdown(f"*{msg}*") | |
# Model and usage information | |
st.sidebar.header("Thông tin mô hình") | |
st.sidebar.info(""" | |
### Mô hình AI Tiếng Việt | |
- Được huấn luyện trên dữ liệu tiếng Việt | |
- Hỗ trợ trò chuyện đa dạng | |
- Lưu ý: Chất lượng trả lời phụ thuộc vào mô hình | |
""") | |
if __name__ == "__main__": | |
main() |