dev/check-gpu-for-rinna

app.py

@@ -25,14 +25,17 @@ E5_EMBEDDINGS = HuggingFaceEmbeddings(
     encode_kwargs=E5_ENCODE_KWARGS,
 )
 
-RINNA_MODEL_NAME = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
-RINNA_TOKENIZER = AutoTokenizer.from_pretrained(RINNA_MODEL_NAME, use_fast=False)
-RINNA_MODEL = AutoModelForCausalLM.from_pretrained(
-    RINNA_MODEL_NAME,
-    load_in_8bit=True,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
+if torch.cuda.is_available():
+    RINNA_MODEL_NAME = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
+    RINNA_TOKENIZER = AutoTokenizer.from_pretrained(RINNA_MODEL_NAME, use_fast=False)
+    RINNA_MODEL = AutoModelForCausalLM.from_pretrained(
+        RINNA_MODEL_NAME,
+        load_in_8bit=True,
+        torch_dtype=torch.float16,
+        device_map="auto",
+    )
+else:
+    RINNA_MODEL = None
 
 
 def _get_config_and_embeddings(collection_name: str | None) -> tuple:
@@ -48,14 +51,17 @@ def _get_config_and_embeddings(collection_name: str | None) -> tuple:
 
 
 def _get_rinna_llm(temperature: float):
-    pipe = pipeline(
-        "text-generation",
-        model=RINNA_MODEL,
-        tokenizer=RINNA_TOKENIZER,
-        max_new_tokens=1024,
-        temperature=temperature,
-    )
-    llm = HuggingFacePipeline(pipeline=pipe)
+    if RINNA_MODEL is not None:
+        pipe = pipeline(
+            "text-generation",
+            model=RINNA_MODEL,
+            tokenizer=RINNA_TOKENIZER,
+            max_new_tokens=1024,
+            temperature=temperature,
+        )
+        llm = HuggingFacePipeline(pipeline=pipe)
+    else:
+        llm = None
     return llm
 
 
@@ -64,7 +70,7 @@ def _get_llm_model(
     temperature: float,
 ):
     if model_name is None:
-        model = "rinna"
+        model = "gpt-3.5-turbo"
     elif model_name == "rinna":
         model = "rinna"
     elif model_name == "GPT-3.5":
@@ -157,12 +163,17 @@ def main(
     return result["result"], html
 
 
+AVAILABLE_LLMS = ["GPT-3.5", "GPT-4"]
+
+if RINNA_MODEL is not None:
+    AVAILABLE_LLMS.append("rinna")
+
 nvdajp_book_qa = gr.Interface(
     fn=main,
     inputs=[
         gr.Textbox(label="query"),
         gr.Radio(["E5", "OpenAI"], label="Embedding", info="選択なしで「E5」を使用"),
-        gr.Radio(["rinna", "GPT-3.5", "GPT-4"], label="Model", info="選択なしで「rinna」を使用"),
+        gr.Radio(AVAILABLE_LLMS, label="Model", info="選択なしで「GPT-3.5」を使用"),
         gr.Radio(
             ["All", "ja-book", "ja-nvda-user-guide", "en-nvda-user-guide"],
             label="絞り込み",