Upload 2 files

app.py

@@ -4,11 +4,18 @@ import os
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from threading import Thread
 import torch
-from torch.nn.attention import SDPBackend, sdpa_kernel
+import gc
+
+def flush():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+torch.set_float32_matmul_precision("high")
 
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 #REPO_ID = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
 REPO_ID = "nicoboss/DeepSeek-R1-Distill-Qwen-32B-Uncensored"
+#REPO_ID = "Qwen/QwQ-32B"
 
 DESCRIPTION = f'''
 <div>
@@ -40,11 +47,13 @@ h1 {
 tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
 if torch.cuda.is_available():
     nf4_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
-    model = AutoModelForCausalLM.from_pretrained(REPO_ID, quantization_config=nf4_config)
+    model = AutoModelForCausalLM.from_pretrained(REPO_ID, device_map="auto", quantization_config=nf4_config)
 else: model = AutoModelForCausalLM.from_pretrained(REPO_ID, torch_dtype=torch.float32)
 streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+flush()
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=59)
+@torch.inference_mode()
 def chat(message: str,
          history: list[dict],
          temperature: float,
@@ -62,7 +71,8 @@ def chat(message: str,
         messages.append({"role": "system", "content": sys_prompt})
         messages.append({"role": "user", "content": message})
 
-        input_tensors = tokenizer.apply_chat_template(history + messages, add_generation_prompt=True, return_dict=True, return_tensors="pt").to(model.device)
+        #input_tensors = tokenizer.apply_chat_template(history + messages, add_generation_prompt=True, return_dict=True, return_tensors="pt").to(model.device)
+        input_tensors = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_dict=True, return_tensors="pt").to(model.device)
 
         input_ids = input_tensors["input_ids"]
         attention_mask = input_tensors["attention_mask"]
@@ -70,8 +80,8 @@ def chat(message: str,
         generate_kwargs = dict(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            streamer=streamer,
             max_new_tokens=max_new_tokens,
+            streamer=streamer,
             do_sample=True,
             temperature=temperature,
             top_k=top_k,
@@ -82,10 +92,8 @@ def chat(message: str,
         if temperature == 0: generate_kwargs['do_sample'] = False
         response.append({"role": "assistant", "content": ""})
 
-        with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
-            t = Thread(target=model.generate, kwargs=generate_kwargs)
-            t.start()
-
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
         for text in streamer:
             response[-1]["content"] += text
             yield response
@@ -93,6 +101,8 @@ def chat(message: str,
         print(e)
         gr.Warning(f"Error: {e}")
         yield response
+    finally:
+        flush()
 
 with gr.Blocks(fill_height=True, fill_width=True, css=css) as demo:
     gr.Markdown(DESCRIPTION)
@@ -108,7 +118,7 @@ with gr.Blocks(fill_height=True, fill_width=True, css=css) as demo:
             gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p", render=False),
             gr.Slider(minimum=0, maximum=100, value=40, step=1, label="Top-k", render=False),
             gr.Slider(minimum=0.0, maximum=2.0, value=1.1, step=0.1, label="Repetition penalty", render=False),
-            gr.Textbox(value="", label="System prompt", render=False),
+            gr.Textbox(value="", label="System prompt", render=False)
         ],
         save_history=True,
         examples=[

requirements.txt

@@ -2,8 +2,10 @@ huggingface_hub
 torch==2.4.0
 torchvision
 accelerate
-transformers
+git+https://github.com/huggingface/transformers
 numpy<2
 sentencepiece
 triton
+optimum
+optimum-quanto
 bitsandbytes