Update app.py

app.py

@@ -10,11 +10,12 @@ from threading import Thread
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = "CohereForAI/aya-23-8B"
 MODEL_ID2 = "CohereForAI/aya-23-35B"
-MODEL_NAME = MODEL_ID2.split("/")[-1]
+MODELS = os.environ.get("MODELS")
+MODEL_NAME = MODELS.split("/")[-1]
 
 TITLE = "<h1><center>Aya-23-Chatbox</center></h1>"
 
-DESCRIPTION = f'<h3><center>MODEL: <a href="https://hf.co/{MODEL_ID}">{MODEL_NAME}</a></center></h3>'
+DESCRIPTION = f'<h3><center>MODEL: <a href="https://hf.co/{MODELS}">{MODEL_NAME}</a></center></h3>'
 
 CSS = """
 .duplicate-button {
@@ -49,13 +50,13 @@ if USE_FLASH_ATTENTION:
     attn_implementation="flash_attention_2"
 
 model = AutoModelForCausalLM.from_pretrained(
-          MODEL_ID2,
+          MODELS,
           quantization_config=quantization_config,
           attn_implementation=attn_implementation,
           torch_dtype=torch.bfloat16,
           device_map="auto",
         )
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID2)
+tokenizer = AutoTokenizer.from_pretrained(MODELS)
 
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):