Update app.py

app.py

@@ -41,20 +41,21 @@ model = Llama4ForConditionalGeneration.from_pretrained(
     # quantization_config=bnb_config,
     device_map="auto",
 )
-
+# processor = AutoProcessor.from_pretrained(model_name, cache_dir = cache_dir)
 tokenizer = AutoTokenizer.from_pretrained(model_name
     # , gguf_file=filename
     # , subfolder=subfolder
     )
-SYSTEM_PROMPT = """
-Respond in the following format:
-<reasoning>
-...
-</reasoning>
-<answer>
-...
-</answer>
-"""
+SYSTEM_PROMPT = "You are a friendly Chatbot."
+# """
+# Respond in the following format:
+# <reasoning>
+# ...
+# </reasoning>
+# <answer>
+# ...
+# </answer>
+# """
 
 @spaces.GPU
 def generate(prompt, history):
@@ -62,25 +63,36 @@ def generate(prompt, history):
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": prompt}
     ]
-    text = tokenizer.apply_chat_template(
+    # text = tokenizer.apply_chat_template(
+    #     messages,
+    #     # tokenize=False,
+    #     tokenize=True,
+    #     add_generation_prompt=True
+    # )
+    # model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    
+    # generated_ids = model.generate(
+    #     **model_inputs,
+    #     max_new_tokens=512
+    # )
+    # generated_ids = [
+    #     output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    # ]
+    
+    # response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    # return response
+    inputs = tokenizer.apply_chat_template(
         messages,
-        # tokenize=False,
+        add_generation_prompt=True,
         tokenize=True,
-        add_generation_prompt=True
-    )
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-    
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=512
+        return_dict=True,
+        return_tensors="pt",
+    ).to(gpu_model.device)
+    outputs = gpu_model.generate(
+        **inputs,
+        max_new_tokens=512,
     )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
-    
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return response
-
+    response = tokenizer.batch_decode(outputs[:, inputs["input_ids"].shape[-1]:])[0]
 
 
 chat_interface = gr.ChatInterface(