.

app.py

@@ -22,22 +22,25 @@ INFERENCE_API = "api"
 # モデル定義
 TEXT_GENERATION_MODELS = [
     {
-        "name": "Zephyr-7B",
-        "description": "Specialized in understanding context and nuance",
+        "name": "Llama-2-7b-chat-hf",
+        "description": "Llama-2-7b-chat-hf",
+        "chat_model": True,
         "type": INFERENCE_API,
-        "model_id": "HuggingFaceH4/zephyr-7b-beta"
+        "model_id": "meta-llama/Llama-2-7b-chat-hf"
     },
     {
-        "name": "Llama-2",
-        "description": "Known for its robust performance in content analysis",
-        "type": LOCAL,
-        "model_path": "meta-llama/Llama-2-7b-hf"
+        "name": "TinyLlaama-1.1B-Chat-v1.0",
+        "description": "TinyLlaama-1.1B-Chat-v1.0",
+        "chat_model": True,
+        "type": INFERENCE_API,
+        "model_id": "tinyllama/TinyLlama-1.1B-Chat-v1.0"
     },
     {
-        "name": "Mistral-7B",
-        "description": "Offers precise and detailed text evaluation",
+        "name": "TinyLlama_v1.1_math_code",
+        "description": "TinyLlama_v1.1_math_code",
+        "chat_model": False,
         "type": LOCAL,
-        "model_path": "mistralai/Mistral-7B-v0.1"
+        "model_path": "TinyLlama/TinyLlama_v1.1_math_code"
     }
 ]
 
@@ -111,7 +114,7 @@ def preload_local_models():
                 logger.error(f"Error preloading model {model_path}: {str(e)}")
 
 @spaces.GPU
-def generate_text_local(model_path, text):
+def generate_text_local(model_path, chat_model, text):
     """ローカルモデルでのテキスト生成"""
     try:
         logger.info(f"Running local text generation with {model_path}")
@@ -129,13 +132,20 @@ def generate_text_local(model_path, text):
         device_info = next(pipeline.model.parameters()).device
         logger.info(f"Model {model_path} is running on device: {device_info}")
         
-        outputs = pipeline(
-            text,
-            max_new_tokens=40,
-            do_sample=False,
-            num_return_sequences=1
-        )
-        
+        if chat_model:
+            outputs = pipeline(
+                [{"role": "user", "content": text}],
+                max_new_tokens=40,
+                do_sample=False,
+                num_return_sequences=1
+            )
+        else:
+            outputs = pipeline(
+                text,
+                max_new_tokens=40,
+                do_sample=False,
+                num_return_sequences=1
+            )
         # モデルをCPUに戻す
         pipeline.model = pipeline.model.to("cpu")
         if hasattr(pipeline, "device"):
@@ -146,15 +156,22 @@ def generate_text_local(model_path, text):
         logger.error(f"Error in local text generation with {model_path}: {str(e)}")
         return f"Error: {str(e)}"
 
-def generate_text_api(model_id, text):
+def generate_text_api(model_id, chat_model, text):
     """API経由でのテキスト生成"""
     try:
         logger.info(f"Running API text generation with {model_id}")
-        response = api_clients[model_id].text_generation(
-            text, 
-            max_new_tokens=40, 
-            temperature=0.7
-        )
+        if chat_model:
+            response = api_clients[model_id].chat.completions.create(
+                messages=[{"role": "user", "content": text}],
+                max_tokens=40,
+                temperature=0.7
+            )
+            response = response.choices[0].message.content
+        else:
+            response = api_clients[model_id].text_generation(
+                text, 
+                max_new_tokens=40, 
+                temperature=0.7)
         return response
     except Exception as e:
         logger.error(f"Error in API text generation with {model_id}: {str(e)}")
@@ -214,11 +231,11 @@ def handle_invoke(text, selected_types):
         for model in TEXT_GENERATION_MODELS:
             if model["type"] in selected_types:
                 if model["type"] == LOCAL:
-                    future = executor.submit(generate_text_local, model["model_path"], text)
+                    future = executor.submit(generate_text_local, model["model_path"], model["chat_model"], text)
                     futures.append(future)
                     futures_to_model[future] = model
                 else:  # api
-                    future = executor.submit(generate_text_api, model["model_id"], text)
+                    future = executor.submit(generate_text_api, model["model_id"], model["chat_model"], text)
                     futures.append(future)
                     futures_to_model[future] = model