Jinja - add missing features for DeepSeek R1 (huggingface#1142)

This PR adds support for the following Jinja features:
* Accessing members of call expression results
* `.split` on strings - the implementation matches [Python's
`str.split(sep=None, maxsplit=-1)` function
behavior](https://docs.python.org/3.13/library/stdtypes.html#str.split)

[DeepSeek R1's chat
template](https://huggingface.co/deepseek-ai/DeepSeek-R1/blob/main/tokenizer_config.json)
used to fail on this expression, which now works in this PR:
```
{%- set content = content.split('</think>')[-1] -%}
```

Fixes huggingface#1141

---------

Co-authored-by: Joshua Lochner <[email protected]>

Author: giladgd

packages/jinja/src/index.ts

@@ -31,7 +31,7 @@ export class Template {
 		this.parsed = parse(tokens);
 	}
 
-	render(items: Record<string, unknown>): string {
+	render(items?: Record<string, unknown>): string {
 		// Create a new environment for this template
 		const env = new Environment();
 
@@ -44,8 +44,10 @@ export class Template {
 		env.set("range", range);
 
 		// Add user-defined variables
-		for (const [key, value] of Object.entries(items)) {
-			env.set(key, value);
+		if (items) {
+			for (const [key, value] of Object.entries(items)) {
+				env.set(key, value);
+			}
 		}
 
 		const interpreter = new Interpreter(env);

packages/jinja/src/parser.ts

@@ -343,7 +343,7 @@ export function parse(tokens: Token[]): Program {
 	function parseCallMemberExpression(): Statement {
 		// Handle member expressions recursively
 
-		const member = parseMemberExpression(); // foo.x
+		const member = parseMemberExpression(parsePrimaryExpression()); // foo.x
 
 		if (is(TOKEN_TYPES.OpenParen)) {
 			// foo.x()
@@ -352,15 +352,17 @@ export function parse(tokens: Token[]): Program {
 		return member;
 	}
 
-	function parseCallExpression(callee: Statement): CallExpression {
-		let callExpression = new CallExpression(callee, parseArgs());
+	function parseCallExpression(callee: Statement): Statement {
+		let expression: Statement = new CallExpression(callee, parseArgs());
+
+		expression = parseMemberExpression(expression); // foo.x().y
 
 		if (is(TOKEN_TYPES.OpenParen)) {
 			// foo.x()()
-			callExpression = parseCallExpression(callExpression);
+			expression = parseCallExpression(expression);
 		}
 
-		return callExpression;
+		return expression;
 	}
 
 	function parseArgs(): Statement[] {
@@ -433,9 +435,7 @@ export function parse(tokens: Token[]): Program {
 		return slices[0] as Statement; // normal member expression
 	}
 
-	function parseMemberExpression(): Statement {
-		let object = parsePrimaryExpression();
-
+	function parseMemberExpression(object: Statement): Statement {
 		while (is(TOKEN_TYPES.Dot) || is(TOKEN_TYPES.OpenSquareBracket)) {
 			const operator = tokens[current]; // . or [
 			++current;

packages/jinja/src/runtime.ts

@@ -117,6 +117,48 @@ export class StringValue extends RuntimeValue<string> {
 				return new StringValue(this.value.trimStart());
 			}),
 		],
+		[
+			"split",
+			// follows Python's `str.split(sep=None, maxsplit=-1)` function behavior
+			// https://docs.python.org/3.13/library/stdtypes.html#str.split
+			new FunctionValue((args) => {
+				const sep = args[0] ?? new NullValue();
+				if (!(sep instanceof StringValue || sep instanceof NullValue)) {
+					throw new Error("sep argument must be a string or null");
+				}
+				const maxsplit = args[1] ?? new NumericValue(-1);
+				if (!(maxsplit instanceof NumericValue)) {
+					throw new Error("maxsplit argument must be a number");
+				}
+
+				let result = [];
+				if (sep instanceof NullValue) {
+					// If sep is not specified or is None, runs of consecutive whitespace are regarded as a single separator, and the
+					// result will contain no empty strings at the start or end if the string has leading or trailing whitespace.
+					// Trailing whitespace may be present when maxsplit is specified and there aren't sufficient matches in the string.
+					const text = this.value.trimStart();
+					for (const { 0: match, index } of text.matchAll(/\S+/g)) {
+						if (maxsplit.value !== -1 && result.length >= maxsplit.value && index !== undefined) {
+							result.push(match + text.slice(index + match.length));
+							break;
+						}
+						result.push(match);
+					}
+				} else {
+					// If sep is specified, consecutive delimiters are not grouped together and are deemed to delimit empty strings.
+					if (sep.value === "") {
+						throw new Error("empty separator");
+					}
+					result = this.value.split(sep.value);
+					if (maxsplit.value !== -1 && result.length > maxsplit.value) {
+						// Follow Python's behavior: If maxsplit is given, at most maxsplit splits are done,
+						// with any remaining text returned as the final element of the list.
+						result.push(result.splice(maxsplit.value).join(sep.value));
+					}
+				}
+				return new ArrayValue(result.map((part) => new StringValue(part)));
+			}),
+		],
 	]);
 }
 
@@ -543,6 +585,8 @@ export class Interpreter {
 								}
 							})
 						);
+					case "join":
+						return new StringValue(operand.value.map((x) => x.value).join(""));
 					default:
 						throw new Error(`Unknown ArrayValue filter: ${filter.value}`);
 				}
@@ -570,6 +614,7 @@ export class Interpreter {
 								)
 								.join("\n")
 						);
+					case "join":
 					case "string":
 						return operand; // no-op
 					default:
@@ -610,6 +655,24 @@ export class Interpreter {
 					throw new Error("If set, indent must be a number");
 				}
 				return new StringValue(toJSON(operand, indent.value));
+			} else if (filterName === "join") {
+				let value;
+				if (operand instanceof StringValue) {
+					// NOTE: string.split('') breaks for unicode characters
+					value = Array.from(operand.value);
+				} else if (operand instanceof ArrayValue) {
+					value = operand.value.map((x) => x.value);
+				} else {
+					throw new Error(`Cannot apply filter "${filterName}" to type: ${operand.type}`);
+				}
+				const [args, kwargs] = this.evaluateArguments(filter.args, environment);
+
+				const separator = args.at(0) ?? kwargs.get("separator") ?? new StringValue("");
+				if (!(separator instanceof StringValue)) {
+					throw new Error("separator must be a string");
+				}
+
+				return new StringValue(value.join(separator.value));
 			}
 
 			if (operand instanceof ArrayValue) {