app.py

import os
from datetime import datetime
import uuid
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch
from huggingface_hub import login
from threading import Thread

from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Get the Hugging Face token from environment variables
hf_token = os.getenv("HUGGINGFACE_TOKEN")

# Load model and tokenizer
model_name = "google/gemma-2-2b-it"

tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
    token=hf_token
)

def chat_with_model(messages):
    # Prepare the input
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Generate response
    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
    generation_kwargs = dict(
        inputs,
        max_new_tokens=1000,
        temperature=0.7,
        do_sample=True,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    return streamer

def chat_with_model_gradio(message, history, session_id):
    system_message = f"너의 이름은 ChatMBTI. 사람들의 MBTI유형에 알맞은 상담을 진행할 수 있어. 상대방의 MBTI 유형을 먼저 물어보고, 그 유형에 알맞게 상담을 진행해줘. 참고로 현재 시각은 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}이야."
    
    messages = [
        # {"role": "system", "content": f"너의 이름은 ChatMBTI. 사람들의 MBTI유형에 알맞은 상담을 진행할 수 있어. 상대방의 MBTI 유형을 먼저 물어보고, 그 유형에 알맞게 상담을 진행해줘. 참고로 현재 시각은 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}이야."},
        {"role": "user", "content": system_message},
        {"role": "assistant", "content": "안녕하세요? ChatMBTI입니다. 오늘 하루 어떠셨나요?"},
    ]
    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for i, (m, _) in enumerate(history)])
    messages.append({"role": "user", "content": message})

    streamer = chat_with_model(messages)
    
    partial_message = ""
    for new_token in streamer:
        partial_message += new_token
        yield "", history + [(message, partial_message)]

def main():
    session_id = str(uuid.uuid4())
    with gr.Blocks() as demo:
        chatbot = gr.Chatbot(label="ChatMBTI")
        msg = gr.Textbox(label="메시지를 입력하세요")
        clear = gr.Button("대화 초기화")

        msg.submit(chat_with_model_gradio, [msg, chatbot, gr.State(session_id)], [msg, chatbot])
        clear.click(lambda: None, None, chatbot, queue=False)

    demo.queue()
    demo.launch()

if __name__ == "__main__":
    main()