Update app.py

app.py

@@ -1,5 +1,4 @@
 import torch
-from PIL import Image
 import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -9,16 +8,16 @@ import random
 from datasets import load_dataset
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL_ID = "aixsatoshi/Meta-Llama-3.1-8B-Instruct-plus-Swallow"
+MODEL_ID = "Aratako/sarashina2.1-1b-sft"
 MODELS = os.environ.get("MODELS")
 MODEL_NAME = MODEL_ID.split("/")[-1]
 
 TITLE = "<h1><center>New japanese LLM model webui</center></h1>"
 
 DESCRIPTION = f"""
-<h3>MODEL: <a href="https://huggingface.co/aixsatoshi/Meta-Llama-3.1-8B-Instruct-plus-Swallow">Meta-Llama-3.1-8B-Instruct-plus-Swallow</a></h3>
+<h3>MODEL: <a href="https://huggingface.co/Aratako/sarashina2.1-1b-sft">Aratako/sarashina2.1-1b-sft</a></h3>
 <center>
-<p>aixsatoshi/Meta-Llama-3.1-8B-Instruct-plus-Swallow is the merged model.
+<p>
 <br>
 Feel free to test without log.
 </p>
@@ -46,7 +45,7 @@ h3 {
 # モデルとトークナイザーの読み込み
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
     device_map="auto",
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -86,7 +85,7 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         max_new_tokens=max_new_tokens, 
         do_sample=True, 
         temperature=temperature,
-        eos_token_id=[128001, 128009],
+        #eos_token_id=[128001, 128009],
     )
     
     thread = Thread(target=model.generate, kwargs=generate_kwargs)