Update app.py

app.py

@@ -7,6 +7,16 @@ from huggingface_hub import InferenceClient
 import os
 import psutil
 
+import json
+import subprocess
+from threading import Thread
+
+import torch
+import spaces
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 
 """
@@ -49,48 +59,48 @@ import transformers
 
 # model_id = "mistralai/Mistral-7B-v0.3"
 
-model_id = "microsoft/Phi-3-medium-4k-instruct"
-# model_id = "microsoft/phi-4"
+# model_id = "microsoft/Phi-3-medium-4k-instruct"
+# # model_id = "microsoft/phi-4"
 
-# model_id = "Qwen/Qwen2-7B-Instruct"
+# # model_id = "Qwen/Qwen2-7B-Instruct"
 
 
-tokenizer = AutoTokenizer.from_pretrained(
-    # model_id
-    model_id,
-    # use_fast=False
-    token= token,
-trust_remote_code=True)
+# tokenizer = AutoTokenizer.from_pretrained(
+#     # model_id
+#     model_id,
+#     # use_fast=False
+#     token= token,
+# trust_remote_code=True)
 
 
-accelerator = Accelerator()
+# accelerator = Accelerator()
 
-model = AutoModelForCausalLM.from_pretrained(model_id, token= token, 
-                                                 # torch_dtype= torch.uint8, 
-                                             torch_dtype=torch.bfloat16,
-                                              # load_in_8bit=True,
-                                             # #  # torch_dtype=torch.fl,
-                                             attn_implementation="flash_attention_2",
-                                             low_cpu_mem_usage=True,
-                                             trust_remote_code=True,
-                                             device_map='cuda',
-                                             # device_map=accelerator.device_map,
+# model = AutoModelForCausalLM.from_pretrained(model_id, token= token, 
+#                                                  # torch_dtype= torch.uint8, 
+#                                              torch_dtype=torch.bfloat16,
+#                                               # load_in_8bit=True,
+#                                              # #  # torch_dtype=torch.fl,
+#                                              attn_implementation="flash_attention_2",
+#                                              low_cpu_mem_usage=True,
+#                                              trust_remote_code=True,
+#                                              device_map='cuda',
+#                                              # device_map=accelerator.device_map,
                                              
-                                            )
+#                                             )
 
 
 
 
 
-# 
-model = accelerator.prepare(model)
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+# # 
+# model = accelerator.prepare(model)
+# from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-)
+# pipe = pipeline(
+#     "text-generation",
+#     model=model,
+#     tokenizer=tokenizer,
+# )
 
 
 
@@ -109,6 +119,27 @@ pipe = pipeline(
 # model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, no_split_module_classes=["GPTJBlock"])
 # model.half()
 
+MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
+CHAT_TEMPLATE = "َAuto"
+MODEL_NAME = MODEL_ID.split("/")[-1]
+CONTEXT_LENGTH = 16000
+
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",
+    quantization_config=quantization_config,
+    attn_implementation="flash_attention_2",
+)
+
+
+
 import json
 
 def str_to_json(str_obj):
@@ -123,48 +154,83 @@ def respond(
     system_message,
     max_tokens,
     temperature,
-    top_p,
-):
+    top_p):
+    
+    stop_tokens = ["<|endoftext|>", "<|im_end|>"]
+    instruction = '<|im_start|>system\n' + system_message + '\n<|im_end|>\n'
+    for user, assistant in history:
+        instruction += f'<|im_start|>user\n{user}\n<|im_end|>\n<|im_start|>assistant\n{assistant}\n<|im_end|>\n'
+    instruction += f'<|im_start|>user\n{message}\n<|im_end|>\n<|im_start|>assistant\n'
+    
+    print(instruction)
+    
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
+    input_ids, attention_mask = enc.input_ids, enc.attention_mask
+
+    if input_ids.shape[1] > CONTEXT_LENGTH:
+        input_ids = input_ids[:, -CONTEXT_LENGTH:]
+        attention_mask = attention_mask[:, -CONTEXT_LENGTH:]
+
+    generate_kwargs = dict(
+        input_ids=input_ids.to(device),
+        attention_mask=attention_mask.to(device),
+        streamer=streamer,
+        do_sample=True,
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        top_p=top_p
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    outputs = []
+    for new_token in streamer:
+        outputs.append(new_token)
+        if new_token in stop_tokens:
+            break
+        yield "".join(outputs)
     # yield 'retuend'
     # model.to(accelerator.device)
 
-    messages = []
-    json_obj = str_to_json(message)
-    print(json_obj)
+#     messages = []
+#     json_obj = str_to_json(message)
+#     print(json_obj)
     
-    messages= json_obj
-
-    # input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(accelerator.device)
-    # input_ids2 = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt") #.to('cuda')
-    # print(f"Converted input_ids dtype: {input_ids.dtype}")
-    # input_str= str(input_ids2)
-    # print('input str = ', input_str)
-
-    generation_args = {
-    "max_new_tokens": max_tokens,
-    "return_full_text": False,
-    "temperature": temperature,
-    "do_sample": False,
-}
-
-    output = pipe(messages, **generation_args)
-    print(output[0]['generated_text'])
-    gen_text=output[0]['generated_text']
-
-    # with torch.no_grad():
-    #     gen_tokens = model.generate(
-    # input_ids, 
-    # max_new_tokens=max_tokens, 
-    # # do_sample=True, 
-    # temperature=temperature,
-    # )
-
-    # gen_text = tokenizer.decode(gen_tokens[0])
-    # print(gen_text)
-    # gen_text= gen_text.replace(input_str,'')
-    # gen_text= gen_text.replace('<|im_end|>','')
+#     messages= json_obj
+
+#     # input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(accelerator.device)
+#     # input_ids2 = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt") #.to('cuda')
+#     # print(f"Converted input_ids dtype: {input_ids.dtype}")
+#     # input_str= str(input_ids2)
+#     # print('input str = ', input_str)
+
+#     generation_args = {
+#     "max_new_tokens": max_tokens,
+#     "return_full_text": False,
+#     "temperature": temperature,
+#     "do_sample": False,
+# }
+
+#     output = pipe(messages, **generation_args)
+#     print(output[0]['generated_text'])
+#     gen_text=output[0]['generated_text']
+
+#     # with torch.no_grad():
+#     #     gen_tokens = model.generate(
+#     # input_ids, 
+#     # max_new_tokens=max_tokens, 
+#     # # do_sample=True, 
+#     # temperature=temperature,
+#     # )
+
+#     # gen_text = tokenizer.decode(gen_tokens[0])
+#     # print(gen_text)
+#     # gen_text= gen_text.replace(input_str,'')
+#     # gen_text= gen_text.replace('<|im_end|>','')
     
-    yield gen_text
+#     yield gen_text
    
   
 #     messages = [