commit

app.py

@@ -10,7 +10,7 @@ Hugging Face Spaces (ZeroGPU) 対応版
 
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import os
 from typing import List, Tuple
 
@@ -34,11 +34,16 @@ class ChatBot:
     def __init__(self):
         self.model = None
         self.tokenizer = None
+        self.pipeline = None
         self.current_model = None
         
+    def is_gpt_oss_model(self, model_name: str) -> bool:
+        """gpt-ossモデルかどうかを判定"""
+        return "gpt-oss" in model_name.lower()
+        
     def load_model(self, model_name: str):
         """モデルとトークナイザーをロード"""
-        if self.current_model == model_name and self.model is not None:
+        if self.current_model == model_name and (self.model is not None or self.pipeline is not None):
             return
             
         try:
@@ -46,47 +51,121 @@ class ChatBot:
             if self.model is not None:
                 del self.model
                 del self.tokenizer
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-                    torch.cuda.synchronize()
-            
-            # トークナイザーロード
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                model_name,
-                token=HF_TOKEN,
-                trust_remote_code=True,
-                padding_side="left"
-            )
-            
-            # パッドトークンの設定
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+            if self.pipeline is not None:
+                del self.pipeline
+                
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                torch.cuda.synchronize()
             
-            # モデルロード（ZeroGPU対応）
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                token=HF_TOKEN,
-                torch_dtype=torch.float16,
-                low_cpu_mem_usage=True,
-                trust_remote_code=True,
-                load_in_8bit=False,  # ZeroGPU環境では8bit量子化は使わない
-                device_map=None  # ZeroGPU環境では自動マッピングしない
-            )
+            if self.is_gpt_oss_model(model_name):
+                # gpt-ossモデルの場合はpipelineを使用
+                print(f"gpt-ossモデル {model_name} をpipelineでロードします...")
+                self.pipeline = pipeline(
+                    "text-generation",
+                    model=model_name,
+                    torch_dtype=torch.float16,
+                    trust_remote_code=True,
+                    token=HF_TOKEN,
+                    device_map=None  # ZeroGPU対応のため手動制御
+                )
+                self.model = None
+                self.tokenizer = None
+            else:
+                # 通常のモデルの場合
+                print(f"通常のモデル {model_name} をロードします...")
+                # トークナイザーロード
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    model_name,
+                    token=HF_TOKEN,
+                    trust_remote_code=True,
+                    padding_side="left"
+                )
+                
+                # パッドトークンの設定
+                if self.tokenizer.pad_token is None:
+                    self.tokenizer.pad_token = self.tokenizer.eos_token
+                    self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+                
+                # モデルロード（ZeroGPU対応）
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    token=HF_TOKEN,
+                    torch_dtype=torch.float16,
+                    low_cpu_mem_usage=True,
+                    trust_remote_code=True,
+                    load_in_8bit=False,  # ZeroGPU環境では8bit量子化は使わない
+                    device_map=None  # ZeroGPU環境では自動マッピングしない
+                )
+                self.pipeline = None
                 
             self.current_model = model_name
             print(f"モデル {model_name} のロードが完了しました。")
             
         except Exception as e:
             print(f"モデルのロード中にエラーが発生しました: {str(e)}")
-            raise
+            # gpt-ossモデルでエラーが出た場合、使用不可と表示
+            if self.is_gpt_oss_model(model_name):
+                raise Exception(f"gpt-ossモデルのロードに失敗しました。このモデルは現在の環境では使用できません: {str(e)}")
+            else:
+                raise
         
     def _generate_response_gpu(self, message: str, history: List[Tuple[str, str]], model_name: str, 
                              temperature: float = 0.7, max_tokens: int = 512) -> str:
         """GPU上で応答を生成する実際の処理"""
-        # モデルロード
-        self.load_model(model_name)
+        try:
+            # モデルロード
+            self.load_model(model_name)
+            
+            if self.is_gpt_oss_model(model_name):
+                # gpt-ossモデルの場合
+                return self._generate_with_pipeline(message, history, temperature, max_tokens)
+            else:
+                # 通常のモデルの場合
+                return self._generate_with_model(message, history, temperature, max_tokens)
+                
+        except Exception as e:
+            return f"エラー: {str(e)}"
+    
+    def _generate_with_pipeline(self, message: str, history: List[Tuple[str, str]], 
+                               temperature: float, max_tokens: int) -> str:
+        """pipelineを使用した生成（gpt-oss用）"""
+        # GPUに移動
+        if hasattr(self.pipeline.model, 'to'):
+            self.pipeline.model.to('cuda')
         
+        # gpt-ossはchat format用のmessages形式を使用
+        messages = []
+        
+        # 履歴を追加（最新3件のみ）
+        for user_msg, assistant_msg in history[-3:]:
+            messages.append({"role": "user", "content": user_msg})
+            messages.append({"role": "assistant", "content": assistant_msg})
+        
+        # 現在のメッセージを追加
+        messages.append({"role": "user", "content": message})
+        
+        # pipeline経由で生成
+        outputs = self.pipeline(
+            messages,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            do_sample=True,
+            top_p=0.95,
+            return_full_text=False
+        )
+        
+        # CPUに戻す（メモリ節約）
+        if hasattr(self.pipeline.model, 'to'):
+            self.pipeline.model.to('cpu')
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+        
+        return outputs[0]["generated_text"].strip()
+    
+    def _generate_with_model(self, message: str, history: List[Tuple[str, str]], 
+                           temperature: float, max_tokens: int) -> str:
+        """通常のモデルを使用した生成"""
         # GPUに移動
         self.model.to('cuda')
         
@@ -128,33 +207,61 @@ class ChatBot:
             return self._generate_response_gpu(message, history, model_name, temperature, max_tokens)
         else:
             # 通常環境の場合
-            self.load_model(model_name)
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
-            
-            if device == 'cuda':
-                self.model.to(device)
-            
-            prompt = self._build_prompt(message, history)
-            inputs = self.tokenizer.encode(prompt, return_tensors="pt").to(device)
-            
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    inputs,
-                    max_new_tokens=max_tokens,
-                    temperature=temperature,
-                    do_sample=True,
-                    top_p=0.95,
-                    top_k=50,
-                    repetition_penalty=1.1,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
-                )
-            
-            response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-            return response.strip()
+            try:
+                self.load_model(model_name)
+                
+                if self.is_gpt_oss_model(model_name):
+                    # gpt-ossモデルの場合
+                    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+                    if hasattr(self.pipeline.model, 'to') and device == 'cuda':
+                        self.pipeline.model.to(device)
+                    
+                    messages = []
+                    for user_msg, assistant_msg in history[-3:]:
+                        messages.append({"role": "user", "content": user_msg})
+                        messages.append({"role": "assistant", "content": assistant_msg})
+                    messages.append({"role": "user", "content": message})
+                    
+                    outputs = self.pipeline(
+                        messages,
+                        max_new_tokens=max_tokens,
+                        temperature=temperature,
+                        do_sample=True,
+                        top_p=0.95,
+                        return_full_text=False
+                    )
+                    return outputs[0]["generated_text"].strip()
+                    
+                else:
+                    # 通常のモデルの場合
+                    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+                    if device == 'cuda':
+                        self.model.to(device)
+                    
+                    prompt = self._build_prompt(message, history)
+                    inputs = self.tokenizer.encode(prompt, return_tensors="pt").to(device)
+                    
+                    with torch.no_grad():
+                        outputs = self.model.generate(
+                            inputs,
+                            max_new_tokens=max_tokens,
+                            temperature=temperature,
+                            do_sample=True,
+                            top_p=0.95,
+                            top_k=50,
+                            repetition_penalty=1.1,
+                            pad_token_id=self.tokenizer.pad_token_id,
+                            eos_token_id=self.tokenizer.eos_token_id
+                        )
+                    
+                    response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+                    return response.strip()
+                    
+            except Exception as e:
+                return f"エラー: {str(e)}"
         
     def _build_prompt(self, message: str, history: List[Tuple[str, str]]) -> str:
-        """会話履歴からプロンプトを構築"""
+        """会話履歴からプロンプトを構築（通常のモデル用）"""
         prompt = ""
         
         # 履歴を追加（最新3件のみ使用 - メモリ効率のため）
@@ -278,7 +385,7 @@ with gr.Blocks(title="ChatGPT Clone", theme=gr.themes.Soft()) as app:
             - ZeroGPU使用により高速推論が可能
             - 1回の生成は120秒以内に完了します
             - 大きなモデル使用時は、短めの応答になる場合があります
-            - gpt-oss-20bは推論専用モデルです
+            - gpt-oss-20bは推論専用モデルで、harmony formatを使用します
             """)
     
     # イベントハンドラ