Update app.py

app.py

@@ -25,21 +25,17 @@ from transformers import (
 from threading import Thread
 
 import gradio as gr
-from dotenv import load_dotenv
 
 import spaces
 
 
-load_dotenv()
-
-HF_API_KEY = os.getenv("HF_API_KEY")
 MODEL_NAME_MAP = {
     "150m-instruct3": "llm-jp/llm-jp-3-150m-instruct3",
     "440m-instruct3": "llm-jp/llm-jp-3-440m-instruct3",
     "980m-instruct3": "llm-jp/llm-jp-3-980m-instruct3",
-    # "1.8b-instruct3": "llm-jp/llm-jp-3-1.8b-instruct3",
-    # "3.7b-instruct3": "llm-jp/llm-jp-3-3.7b-instruct3",
-    # "13b-instruct3": "llm-jp/llm-jp-3-13b-instruct3",
+    "1.8b-instruct3": "llm-jp/llm-jp-3-1.8b-instruct3",
+    "3.7b-instruct3": "llm-jp/llm-jp-3-3.7b-instruct3",
+    "13b-instruct3": "llm-jp/llm-jp-3-13b-instruct3",
 }
 
 quantization_config = BitsAndBytesConfig(
@@ -50,12 +46,14 @@ quantization_config = BitsAndBytesConfig(
 )
 MODELS = {
     key: AutoModelForCausalLM.from_pretrained(
-        repo_id, quantization_config=quantization_config, device_map="auto"
+        repo_id, 
+        quantization_config=quantization_config, 
+        device_map="auto",
+        attn_implementation="flash_attention_2",
     ) for key, repo_id in MODEL_NAME_MAP.items()
 }
 TOKENIZERS = {
     key: AutoTokenizer.from_pretrained(repo_id) for key, repo_id in MODEL_NAME_MAP.items()
-    
 }
 
 print("Compiling model...")