Update app.py

app.py

@@ -2,18 +2,18 @@ import torch
 from PIL import Image
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import os
 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_ID = "THUDM/glm-4-9b-chat"
+MODEL_ID2 = "THUDM/glm-4-9b-chat-1m"
 MODELS = os.environ.get("MODELS")
 MODEL_NAME = MODELS.split("/")[-1]
 
-TITLE = "<h1><center>Aya-23-Chatbox</center></h1>"
+TITLE = "<h1><center>GLM-4-9B</center></h1>"
 
 DESCRIPTION = f'<h3><center>MODEL: <a href="https://hf.co/{MODELS}">{MODEL_NAME}</a></center></h3>'
 
@@ -26,37 +26,14 @@ CSS = """
 }
 """
 
-
-#QUANTIZE
-QUANTIZE_4BIT = True
-USE_GRAD_CHECKPOINTING = True
-TRAIN_BATCH_SIZE = 2
-TRAIN_MAX_SEQ_LENGTH = 512
-USE_FLASH_ATTENTION = False
-GRAD_ACC_STEPS = 16
-
-quantization_config = None
-
-if QUANTIZE_4BIT:
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_compute_dtype=torch.bfloat16,
-    )
-
-attn_implementation = None
-if USE_FLASH_ATTENTION:
-    attn_implementation="flash_attention_2"
-
 model = AutoModelForCausalLM.from_pretrained(
-          MODELS,
-          quantization_config=quantization_config,
-          attn_implementation=attn_implementation,
-          torch_dtype=torch.bfloat16,
-          device_map="auto",
-        )
-tokenizer = AutoTokenizer.from_pretrained(MODELS)
+        MODELS,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True,
+        ).to(0).eval()
+
+tokenizer = AutoTokenizer.from_pretrained(MODELS,trust_remote_code=True)
 
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):
@@ -69,7 +46,7 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
 
     print(f"Conversation is -\n{conversation}")
     
-    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
+    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
     
     streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": True, "skip_prompt": True, 'clean_up_tokenization_spaces':False,})
 
@@ -79,6 +56,8 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         max_new_tokens=max_new_tokens, 
         do_sample=True, 
         temperature=temperature,
+        repetition_penalty=1.2,
+        eos_token_id=model.config.eos_token_id,
     )
     
     thread = Thread(target=model.generate, kwargs=generate_kwargs)