Update core/make_pipeline.py

core/make_pipeline.py

@@ -1,11 +1,19 @@
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
-import spaces
+
+try:
+    import spaces
+except ImportError:
+    class DummyDecorator:
+        @staticmethod
+        def GPU(func):
+            return func
+    spaces = DummyDecorator()
 
 class MakePipeline:
     # 모델명
-    MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"
+    MODEL_ID = "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B"
     
     # 변수초기화
     # model_id
@@ -24,7 +32,7 @@ class MakePipeline:
             "repetition_penalty": 1.05,
             "max_new_tokens": 96
         }
-    
+
     # 모델 불러오기
     @spaces.GPU
     def build(self, type: str):
@@ -50,7 +58,6 @@ class MakePipeline:
                 "text-generation",
                 model=model,
                 tokenizer=tokenizer,
-                torch_dtype=torch.float16
             )
 
         else:
@@ -82,6 +89,7 @@ class MakePipeline:
             temperature=self.config["temperature"],
             top_p=self.config["top_p"],
             repetition_penalty=self.config["repetition_penalty"],
+            eos_token_id=self.tokenizer.eos_token_id,
             return_full_text=True
         )
         return outputs[0]["generated_text"]