Upload 2 files

app.py

@@ -15,6 +15,7 @@ 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 = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 
 DESCRIPTION = f'''
 <div>
@@ -49,11 +50,10 @@ if torch.cuda.is_available():
     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=59)
 @torch.inference_mode()
-def chat(message: str,
+def chat_stream(message: str,
          history: list[dict],
          temperature: float,
          max_new_tokens: int,
@@ -70,11 +70,15 @@ def chat(message: str,
         messages.append({"role": "system", "content": sys_prompt})
         messages.append({"role": "user", "content": message})
 
-        input_tensors = tokenizer.apply_chat_template([{"role": x["role"], "content": x["content"]} for x in history] + messages, add_generation_prompt=True, return_dict=True, add_special_tokens=False, return_tensors="pt").to(model.device)
+        input_tensors = tokenizer.apply_chat_template([{"role": x["role"], "content": x["content"]} for x in history + messages if "role" in x.keys()], add_generation_prompt=True, return_dict=True, add_special_tokens=False, return_tensors="pt").to(model.device)
 
         input_ids = input_tensors["input_ids"]
         attention_mask = input_tensors["attention_mask"]
 
+        #print("history: ", [{"role": x["role"], "content": x["content"]} for x in history if "role" in x.keys()])
+        #print("messages: ", [{"role": x["role"], "content": x["content"]} for x in messages if "role" in x.keys()])
+        #print("tokenized: ", tokenizer.apply_chat_template([{"role": x["role"], "content": x["content"]} for x in history + messages if "role" in x.keys()], add_generation_prompt=True, add_special_tokens=False, tokenize=False))
+
         generate_kwargs = dict(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -102,10 +106,60 @@ def chat(message: str,
     finally:
         flush()
 
+@spaces.GPU(duration=59)
+@torch.inference_mode()
+def chat(message: str,
+         history: list[dict],
+         temperature: float,
+         max_new_tokens: int,
+         top_p: float,
+         top_k: int,
+         repetition_penalty: float,
+         sys_prompt: str,
+         progress=gr.Progress(track_tqdm=True)
+        ):
+    try:
+        messages = []
+        response = []
+        if not history: history = []
+        messages.append({"role": "system", "content": sys_prompt})
+        messages.append({"role": "user", "content": message})
+
+        input_tensors = tokenizer.apply_chat_template([{"role": x["role"], "content": x["content"]} for x in history + messages if "role" in x.keys()], add_generation_prompt=True, return_dict=True, add_special_tokens=False, return_tensors="pt").to(model.device)
+
+        input_ids = input_tensors["input_ids"]
+        attention_mask = input_tensors["attention_mask"]
+
+        generate_kwargs = dict(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+        if temperature == 0: generate_kwargs['do_sample'] = False
+        response.append({"role": "assistant", "content": ""})
+
+        output_ids = model.generate(**generate_kwargs)
+        output = tokenizer.decode(output_ids.tolist()[0][input_ids.size(1) :], skip_special_tokens=True)
+
+        response[-1]["content"] = output
+        return response
+    except Exception as e:
+        print(e)
+        gr.Warning(f"Error: {e}")
+        return response
+    finally:
+        flush()
+
 with gr.Blocks(fill_height=True, fill_width=True, css=css) as demo:
     gr.Markdown(DESCRIPTION)
     gr.ChatInterface(
-        fn=chat,
+        fn=chat_stream,
         type="messages",
         chatbot=gr.Chatbot(height=450, type="messages", placeholder=PLACEHOLDER, label='Gradio ChatInterface'),
         fill_height=True,