which model would work?

app.py

@@ -2,67 +2,58 @@ import os
 from datetime import datetime
 import uuid
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 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")
+# Authenticate with Hugging Face
+login(token=os.getenv("HUGGINGFACE_TOKEN"))
 
 # Load model and tokenizer
-model_name = "google/gemma-2-2b-it"
+model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", token=True)
+
 
-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
-)
+# Set pad_token_id if it's not already set
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
 
 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)
+    input_ids = tokenizer.encode(str(messages), return_tensors="pt").to(model.device)
+    attention_mask = torch.ones_like(input_ids)
     
     # 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
+    with torch.no_grad():
+        output = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            max_length=1000,
+            num_return_sequences=1,
+            temperature=0.7,
+            pad_token_id=tokenizer.pad_token_id
+        )
+    
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
 
 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입니다. 오늘 하루 어떠셨나요?"},
+        {"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 "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)]
+    response = chat_with_model(messages)
+    history.append((message, response))
+
+    return "", history
 
 def main():
     session_id = str(uuid.uuid4())
@@ -74,8 +65,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__":
-    main()
+    main()