Upload app.py

app.py

@@ -7,26 +7,22 @@ import spaces
 import torch
 
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 ACCESS_TOKEN = os.getenv("HF_TOKEN", "")
 
-model_id = "Qwen/Qwen2.5-14B-Instruct-GPTQ-Int8"
-
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+model_id = "Qwen/Qwen2.5-7B-Instruct"
 
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.float16,
-    device_map="cuda",
+    device_map="auto",
     trust_remote_code=True,
     low_cpu_mem_usage=True,
-    quantize_config=quantization_config,
-    token=ACCESS_TOKEN).to("cuda")
-
+    token=ACCESS_TOKEN)
 tokenizer = AutoTokenizer.from_pretrained(
     model_id,
     trust_remote_code=True,
@@ -47,7 +43,7 @@ def generate(
         conversation.append({"role": "system", "content": system_prompt})
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to("cuda")
+    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")