Update app.py

app.py

@@ -76,8 +76,18 @@ model = AutoModelForCausalLM.from_pretrained(model_id, token= token,
                                             )
 
 
+
+
+
 # 
 model = accelerator.prepare(model)
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+)
 
 
 # device_map = infer_auto_device_map(model, max_memory={0: "79GB", "cpu":"65GB" })
@@ -111,24 +121,35 @@ def respond(
     
     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)
-
-    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|>','')
+    # 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