.

app.py

@@ -61,8 +61,6 @@ CLASSIFICATION_MODELS = [
 
 class LocalModelManager:
     def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"Using device: {self.device}")
         self.models = {}
         self.tokenizers = {}
         self.pipelines = {}
@@ -75,42 +73,36 @@ class LocalModelManager:
                 self.tokenizers[model_path] = AutoTokenizer.from_pretrained(model_path)
                 
                 if task == "text-generation":
-                    model = AutoModelForCausalLM.from_pretrained(
-                        model_path,
-                        torch_dtype=torch.float16,
-                        device_map="auto"
-                    )
                     self.pipelines[model_path] = pipeline(
                         "text-generation",
-                        model=model,
-                        tokenizer=self.tokenizers[model_path]
-                    )
-                else:  # classification
-                    model = AutoModelForSequenceClassification.from_pretrained(
-                        model_path,
+                        model=model_path,
+                        tokenizer=self.tokenizers[model_path],
+                        torch_dtype=torch.bfloat16,
+                        trust_remote_code=True,
                         device_map="auto"
                     )
+                else:  # classification
                     self.pipelines[model_path] = pipeline(
                         "text-classification",
-                        model=model,
-                        tokenizer=self.tokenizers[model_path]
+                        model=model_path,
+                        tokenizer=self.tokenizers[model_path],
+                        torch_dtype=torch.bfloat16,
+                        trust_remote_code=True,
+                        device_map="auto"
                     )
                 
-                self.models[model_path] = model
                 logger.info(f"Model loaded successfully: {model_path}")
             except Exception as e:
                 logger.error(f"Error loading model {model_path}: {str(e)}")
                 raise
 
-    @spaces.GPU()
+    @spaces.GPU
     def _generate_text_sync(self, pipeline, text: str) -> str:
         """同期的なテキスト生成の実行"""
         outputs = pipeline(
             text,
             max_new_tokens=100,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
+            do_sample=False,
             num_return_sequences=1
         )
         return outputs[0]["generated_text"]
@@ -126,7 +118,7 @@ class LocalModelManager:
             logger.error(f"Error in text generation with {model_path}: {str(e)}")
             raise
 
-    @spaces.GPU()
+    @spaces.GPU
     def _classify_text_sync(self, pipeline, text: str) -> str:
         """同期的なテキスト分類の実行"""
         result = pipeline(text)