Update app.py

app.py

@@ -9,18 +9,10 @@ from tokenization_yi import YiTokenizer
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:50'
 model_id = "TheBloke/Yi-34B-200K-Llamafied-GPTQ"
 
-gptq_config = GPTQConfig(
-    bits=4, 
-    exllama_config={"version": 2},
-    disable_exllama=True
-)
+gptq_config = GPTQConfig( bits=4, exllama_config={"version": 2})
 tokenizer = YiTokenizer.from_pretrained("./")
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    device_map="auto", 
-    quantization_config=gptq_config
+model = AutoModelForCausalLM.from_pretrained( model_id, device_map="cuda", 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')