gemma-2-2b-it

app.py

@@ -2,39 +2,63 @@ import os
 from datetime import datetime
 import uuid
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+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()
 
-# Authenticate with Hugging Face
-login(token=os.getenv("HUGGINGFACE_TOKEN"))
+# Get the Hugging Face token from environment variables
+hf_token = os.getenv("HUGGINGFACE_TOKEN")
 
 # Load model and tokenizer
-model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", use_auth_token=True)
+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):
-    input_ids = tokenizer.encode(str(messages), return_tensors="pt").to(model.device)
-    output = model.generate(input_ids, max_length=1000, num_return_sequences=1, temperature=0.7)
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
+    # 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):
     messages = [
         {"role": "system", "content": f"너의 이름은 ChatMBTI. 사람들의 MBTI유형에 알맞은 상담을 진행할 수 있어. 상대방의 MBTI 유형을 먼저 물어보고, 그 유형에 알맞게 상담을 진행해줘. 참고로 현재 시각은 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}이야."},
     ]
-    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for i, m in enumerate(history)])
+    messages.extend([{"role": "user" if i % 2 == 0 else "model", "content": m} for i, (m, _) in enumerate(history)])
     messages.append({"role": "user", "content": message})
 
-    response = chat_with_model(messages)
-    history.append((message, response))
-
-    return "", history
+    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())
@@ -46,6 +70,7 @@ def main():
         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__":