Upload folder using huggingface_hub

backend.py

@@ -88,15 +88,18 @@ class Inference:
         tokenizer = self.models[model_id]["tokenizer"]
         model = self.models[model_id]["model"]
 
-        conversation = ""
+        messages = []
         for msg in history:
             role = msg.get("role", "user")
             content = msg.get("content", "")
-            if role == "user":
-                conversation += f"User: {content}\n"
-            else:
-                conversation += f"Assistant: {content}\n"
-        conversation += f"User: {message}\nAssistant:"
+            messages.append({"role": role, "content": content})
+        messages.append({"role": "user", "content": message})
+
+        conversation = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
 
         try:
             inputs = tokenizer(conversation, return_tensors="pt").to("cuda")
@@ -114,8 +117,9 @@ class Inference:
                 )
             logger.info(f"Generated output shape: {outputs.shape}")
 
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            response = response.split("Assistant:")[-1].strip()
+            # Extract only the newly generated tokens (skip the input)
+            new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+            response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
             logger.info(f"Final response length: {len(response)}")
             logger.info(f"Response: {response}")