Update app.py

app.py

@@ -1,23 +1,22 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, GPTQConfig
 from tokenization_yi import YiTokenizer
 import torch
 import os
 import gradio as gr
 import sentencepiece
 
-model_id = "01-ai/Yi-34B-200K"
-
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
-device = "cuda" if torch.cuda.is_available() else "cpu"
-# offload_directory = './model_offload'
-# if not os.path.exists(offload_directory):
-#     os.makedirs(offload_directory)
-# model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True, load_in_8bit_fp32_cpu_offload=True, offload_folder=offload_directory, trust_remote_code=True)
-# model = model.to(device)
-    
-tokenizer = YiTokenizer(vocab_file="./tokenizer.model")
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True, trust_remote_code=True)
-
+model_id = "TheBloke/Yi-34B-200K-Llamafied-GPTQ"
+
+gptq_config = GPTQConfig(
+    bits=4, 
+    exllama_config={"version": 2}
+)
+tokenizer = AutoTokenizer.from_pretrained("TheBloke/Yi-34B-200K-Llamafied-GPTQ")
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    device_map="auto", 
+    quantization_config=gptq_config
+)
 def run(message, chat_history, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800):
     prompt = get_prompt(message, chat_history)
     input_ids = tokenizer.encode(prompt, return_tensors='pt')