Update app.py

app.py

@@ -1,36 +1,10 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from mergekit.config import MergeConfig
-from mergekit.merge import Merger
 
-# MergeKit の設定を読み込む
-merge_config = MergeConfig.from_dict({
-    "slices": [
-        {
-            "sources": [
-                {"model": "Qwen/Qwen2.5-0.5B-Instruct", "layer_range": [0, 23]},
-                {"model": "Qwen/Qwen2.5-0.5B-Instruct", "layer_range": [0, 23]}
-            ]
-        }
-    ],
-    "merge_method": "slerp",
-    "base_model": "Qwen/Qwen2.5-1.5B-Instruct",
-    "parameters": {
-        "t": [
-            {"filter": "self_attn", "value": [0, 0.5, 0.3, 0.7, 1]},
-            {"filter": "mlp", "value": [1, 0.5, 0.7, 0.3, 0]},
-            {"value": 0.5}
-        ]
-    },
-    "dtype": "bfloat16"
-})
-
-# マージされたモデルを生成
-merger = Merger(merge_config)
-merged_model = merger.merge()
-
-# トークナイザーの読み込み
-tokenizer = AutoTokenizer.from_pretrained(merge_config.base_model)
+# モデルとトークナイザーの読み込み
+model_name = "Sakalti/chromet1-0.48b"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, ignore_mismatched_sizes=True)
 
 # 生成用の関数
 def respond(input_text, system_message, max_new_tokens, temperature, top_p):
@@ -44,7 +18,7 @@ def respond(input_text, system_message, max_new_tokens, temperature, top_p):
     print(inputs['input_ids'].shape)
     
     # モデルに入力を渡して生成
-    outputs = merged_model.generate(
+    outputs = model.generate(
         input_ids=inputs['input_ids'],
         attention_mask=inputs['attention_mask'],
         max_new_tokens=max_new_tokens,  # 新規トークンの最大数
@@ -62,7 +36,7 @@ def respond(input_text, system_message, max_new_tokens, temperature, top_p):
 
 # Gradioインターフェースの作成
 with gr.Blocks() as demo:
-    gr.Markdown("## Sakalti/aquamarine チャットボット")
+    gr.Markdown("## chromet チャットボット")
 
     # 追加の入力フィールドをリストで設定
     additional_inputs = [