Update app.py

app.py

@@ -17,8 +17,10 @@ HF_TOKEN = os.environ.get("HF_TOKEN", None)
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
 
-model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-peft_model_id = "Imran1/Llama3.1_8b_Qlora_bnk"
+# model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+# peft_model_id = "Imran1/Llama3.1_8b_Qlora_bnk"
+model_id = "Qwen/Qwen2.5-14B-Instruct"
+peft_model_id = "Imran1/Qwen2.5-14b-bnk-lora-11"
  #attn_implementation="flash_attention_2",
 model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="sdpa",  torch_dtype= torch.bfloat16)
 model.load_adapter(peft_model_id)
@@ -32,10 +34,10 @@ if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
 # Define terminators
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
+# terminators = [
+#     tokenizer.eos_token_id,
+#     tokenizer.convert_tokens_to_ids("<|eot_id|>")
+# ]
 
 generation_params = {
     'max_new_tokens': 2000,
@@ -44,8 +46,8 @@ generation_params = {
     'temperature': 0.7,
     'top_p': 0.9,
     # 'top_k': 50,
-    'pad_token_id': tokenizer.pad_token_id,
-    'eos_token_id': terminators,
+    # 'pad_token_id': tokenizer.pad_token_id,
+    # 'eos_token_id': terminators,
 }
 
 
@@ -89,9 +91,10 @@ def inference(query):
     tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to("cuda")
     outputs = model.generate(tokenized_chat, **generation_params)
     decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=False)
-    assistant_response = decoded_outputs[0].split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
-    response_ = assistant_response.replace('<|eot_id|>', "")
-    return assistant_response
+    # assistant_response = decoded_outputs[0].split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
+    response = decoded_outputs[0][tokenized_chat.shape[-1]:]
+    response = tokenizer.decode(response, skip_special_tokens=True))
+    return response
     # outputs = model.generate(tokenized_chat, **generation_params, streamer=streamer)
     # return outputs