Fix bugs: use token param, apply Llama 3.1 chat template, decode only new tokens

app.py

@@ -37,8 +37,8 @@ def load_model_and_tokenizer():
 
     if USE_PEFT:
         try:
-            _ = PeftConfig.from_pretrained(PEFT_MODEL_ID, use_auth_token=HF_TOKEN)
-            model = PeftModel.from_pretrained(base, PEFT_MODEL_ID, use_auth_token=HF_TOKEN)
+            _ = PeftConfig.from_pretrained(PEFT_MODEL_ID, token=HF_TOKEN)
+            model = PeftModel.from_pretrained(base, PEFT_MODEL_ID, token=HF_TOKEN)
             print(f"[INFO] Loaded PEFT adapter: {PEFT_MODEL_ID}")
             return model, tok
         except Exception as e:
@@ -51,9 +51,17 @@ DEVICE = model.device
 
 @spaces.GPU
 @torch.inference_mode()
-def generate_response(prompt: str, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+def generate_response(messages, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+    # Apply Llama 3.1 chat template
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
     enc = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
     enc = {k: v.to(DEVICE) for k, v in enc.items()}
+
+    input_length = enc['input_ids'].shape[1]
     out = model.generate(
         **enc,
         max_new_tokens=max_new_tokens,
@@ -62,30 +70,30 @@ def generate_response(prompt: str, max_new_tokens=512, temperature=0.7, top_p=0.
         top_p=top_p,
         pad_token_id=tokenizer.eos_token_id,
     )
-    return tokenizer.decode(out[0], skip_special_tokens=True)
+    # Decode only the newly generated tokens
+    return tokenizer.decode(out[0][input_length:], skip_special_tokens=True)
 
 def chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p):
-    # Build a simple conversation string
-    conv = []
+    # Build conversation in Llama 3.1 chat format
+    messages = []
     if system_prompt:
-        conv.append(f"system: {system_prompt}")
-    for u, a in (history or []):
-        if u:
-            conv.append(f"user: {u}")
-        if a:
-            conv.append(f"assistant: {a}")
+        messages.append({"role": "system", "content": system_prompt})
+
+    for user_msg, assistant_msg in (history or []):
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
+
     if message:
-        conv.append(f"user: {message}")
-    prompt = "\n".join(conv) + "\nassistant:"
+        messages.append({"role": "user", "content": message})
+
     reply = generate_response(
-        prompt,
+        messages,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         top_p=top_p,
     )
-    # Strip trailing
-    if "assistant:" in reply:
-        reply = reply.split("assistant:")[-1].strip()
     return reply
 
 demo = gr.ChatInterface(