Skip to content

Commit 55b4ab5

Browse files
SBrandeisalexrs-cohere
SBrandeis
authored and
alexrs-cohere
committed
factor + tapes.json
1 parent 429e81f commit 55b4ab5

File tree

4 files changed

+91
-117
lines changed

4 files changed

+91
-117
lines changed

‎packages/inference/src/tasks/nlp/chatCompletion.ts‎

Lines changed: 11 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,26 @@
11
import { InferenceOutputError } from "../../lib/InferenceOutputError";
2+
import type { CohereTextGenerationOutputFinishReason, CohereMessage, CohereLogprob } from "../../providers/cohere";
23
import type { BaseArgs, Options } from "../../types";
34
import { request } from "../custom/request";
45
import type { ChatCompletionInput, ChatCompletionOutput } from "@huggingface/tasks";
56

6-
export type CohereTextGenerationOutputFinishReason =
7-
| "COMPLETE"
8-
| "STOP_SEQUENCE"
9-
| "MAX_TOKENS"
10-
| "TOOL_CALL"
11-
| "ERROR";
12-
137
interface CohereChatCompletionOutput {
148
id: string;
159
finish_reason: CohereTextGenerationOutputFinishReason;
1610
message: CohereMessage;
17-
usage: CohereChatCompletionOutputUsage;
11+
usage: {
12+
billed_units: {
13+
input_tokens: number;
14+
output_tokens: number;
15+
};
16+
tokens: {
17+
input_tokens: number;
18+
output_tokens: number;
19+
};
20+
};
1821
logprobs?: CohereLogprob[]; // Optional field for log probabilities
1922
}
2023

21-
interface CohereMessage {
22-
role: string;
23-
content: Array<{
24-
type: string;
25-
text: string;
26-
}>;
27-
tool_calls?: CohereToolCall[]; // Optional field for tool calls
28-
}
29-
30-
interface CohereChatCompletionOutputUsage {
31-
billed_units: CohereInputOutputTokens;
32-
tokens: CohereInputOutputTokens;
33-
}
34-
35-
interface CohereInputOutputTokens {
36-
input_tokens: number;
37-
output_tokens: number;
38-
}
39-
40-
interface CohereLogprob {
41-
logprob: number;
42-
token: string;
43-
top_logprobs: CohereTopLogprob[];
44-
}
45-
46-
interface CohereTopLogprob {
47-
logprob: number;
48-
token: string;
49-
}
50-
51-
interface CohereToolCall {
52-
function: CohereFunctionDefinition;
53-
id: string;
54-
type: string;
55-
}
56-
57-
interface CohereFunctionDefinition {
58-
arguments: unknown;
59-
description?: string;
60-
name: string;
61-
}
62-
6324
function convertCohereToChatCompletionOutput(res: CohereChatCompletionOutput): ChatCompletionOutput {
6425
// Create a ChatCompletionOutput object from the CohereChatCompletionOutput
6526
return {

‎packages/inference/src/tasks/nlp/chatCompletionStream.ts‎

Lines changed: 14 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,66 +1,25 @@
1+
import type { CohereLogprob, CohereMessageDelta, CohereTextGenerationOutputFinishReason } from "../../providers/cohere";
12
import type { BaseArgs, Options } from "../../types";
23
import { streamingRequest } from "../custom/streamingRequest";
34
import type { ChatCompletionInput, ChatCompletionStreamOutput } from "@huggingface/tasks";
45

5-
export type CohereTextGenerationOutputFinishReason =
6-
| "COMPLETE"
7-
| "STOP_SEQUENCE"
8-
| "MAX_TOKENS"
9-
| "TOOL_CALL"
10-
| "ERROR";
11-
126
interface CohereChatCompletionStreamOutput {
137
id: string;
148
finish_reason?: CohereTextGenerationOutputFinishReason;
15-
delta: CohereMessageDelta;
16-
usage?: CohereChatCompletionOutputUsage;
17-
logprobs?: CohereLogprob[];
18-
}
19-
20-
interface CohereMessage {
21-
role: string;
22-
content: {
23-
type: string;
24-
text: string;
9+
delta: {
10+
message: CohereMessageDelta;
2511
};
26-
tool_calls?: CohereToolCall[];
27-
}
28-
29-
interface CohereMessageDelta {
30-
message: CohereMessage;
31-
}
32-
33-
interface CohereChatCompletionOutputUsage {
34-
billed_units: CohereInputOutputTokens;
35-
tokens: CohereInputOutputTokens;
36-
}
37-
38-
interface CohereInputOutputTokens {
39-
input_tokens: number;
40-
output_tokens: number;
41-
}
42-
43-
interface CohereLogprob {
44-
logprob: number;
45-
token: string;
46-
top_logprobs: CohereTopLogprob[];
47-
}
48-
49-
interface CohereTopLogprob {
50-
logprob: number;
51-
token: string;
52-
}
53-
54-
interface CohereToolCall {
55-
function: CohereFunctionDefinition;
56-
id: string;
57-
type: string;
58-
}
59-
60-
interface CohereFunctionDefinition {
61-
arguments: unknown;
62-
description?: string;
63-
name: string;
12+
usage?: {
13+
billed_units: {
14+
input_tokens: number;
15+
output_tokens: number;
16+
};
17+
tokens: {
18+
input_tokens: number;
19+
output_tokens: number;
20+
};
21+
};
22+
logprobs?: CohereLogprob[];
6423
}
6524

6625
function convertCohereToChatCompletionStreamOutput(res: CohereChatCompletionStreamOutput): ChatCompletionStreamOutput {

‎packages/inference/test/HfInference.spec.ts‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1355,21 +1355,22 @@ describe.concurrent("HfInference", () => {
13551355
() => {
13561356
const client = new HfInference(env.HF_COHERE_KEY);
13571357

1358-
(HARDCODED_MODEL_ID_MAPPING["cohere"] = {
1358+
HARDCODED_MODEL_ID_MAPPING["cohere"] = {
13591359
"CohereForAI/c4ai-command-r7b-12-2024": "command-r7b-12-2024",
13601360
"CohereForAI/aya-expanse-8b": "c4ai-aya-expanse-8b",
1361-
}),
1362-
it("chatCompletion", async () => {
1363-
const res = await client.chatCompletion({
1364-
model: "CohereForAI/c4ai-command-r7b-12-2024",
1365-
provider: "cohere",
1366-
messages: [{ role: "user", content: "Complete this sentence with words, one plus one is equal " }],
1367-
});
1368-
if (res.choices && res.choices.length > 0) {
1369-
const completion = res.choices[0].message?.content;
1370-
expect(completion).toContain("two");
1371-
}
1361+
};
1362+
1363+
it("chatCompletion", async () => {
1364+
const res = await client.chatCompletion({
1365+
model: "CohereForAI/c4ai-command-r7b-12-2024",
1366+
provider: "cohere",
1367+
messages: [{ role: "user", content: "Complete this sentence with words, one plus one is equal " }],
13721368
});
1369+
if (res.choices && res.choices.length > 0) {
1370+
const completion = res.choices[0].message?.content;
1371+
expect(completion).toContain("two");
1372+
}
1373+
});
13731374

13741375
it("chatCompletion stream", async () => {
13751376
const stream = client.chatCompletionStream({

‎packages/inference/test/tapes.json‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7386,5 +7386,58 @@
73867386
"content-type": "image/jpeg"
73877387
}
73887388
}
7389+
},
7390+
"772e481d98640490fca3aab8e7fed5b771ea213f2b76b2f8858ce7bc90acb16b": {
7391+
"url": "https://api.cohere.com/v2/chat",
7392+
"init": {
7393+
"headers": {
7394+
"Content-Type": "application/json"
7395+
},
7396+
"method": "POST",
7397+
"body": "{\"messages\":[{\"role\":\"user\",\"content\":\"Say 'this is a test'\"}],\"stream\":true,\"model\":\"command-r7b-12-2024\"}"
7398+
},
7399+
"response": {
7400+
"body": "event: message-start\ndata: {\"id\":\"b9c2d3f2-4532-473d-a8a1-9236d159ab26\",\"type\":\"message-start\",\"delta\":{\"message\":{\"role\":\"assistant\",\"content\":[],\"tool_plan\":\"\",\"tool_calls\":[],\"citations\":[]}}}\n\nevent: content-start\ndata: {\"type\":\"content-start\",\"index\":0,\"delta\":{\"message\":{\"content\":{\"type\":\"text\",\"text\":\"\"}}}}\n\nevent: content-delta\ndata: {\"type\":\"content-delta\",\"index\":0,\"delta\":{\"message\":{\"content\":{\"text\":\"This\"}}}}\n\nevent: content-delta\ndata: {\"type\":\"content-delta\",\"index\":0,\"delta\":{\"message\":{\"content\":{\"text\":\" is\"}}}}\n\nevent: content-delta\ndata: {\"type\":\"content-delta\",\"index\":0,\"delta\":{\"message\":{\"content\":{\"text\":\" a\"}}}}\n\nevent: content-delta\ndata: {\"type\":\"content-delta\",\"index\":0,\"delta\":{\"message\":{\"content\":{\"text\":\" test\"}}}}\n\nevent: content-delta\ndata: {\"type\":\"content-delta\",\"index\":0,\"delta\":{\"message\":{\"content\":{\"text\":\".\"}}}}\n\nevent: content-end\ndata: {\"type\":\"content-end\",\"index\":0}\n\nevent: message-end\ndata: {\"type\":\"message-end\",\"delta\":{\"finish_reason\":\"COMPLETE\",\"usage\":{\"billed_units\":{\"input_tokens\":7,\"output_tokens\":5},\"tokens\":{\"input_tokens\":502,\"output_tokens\":7}}}}\n\ndata: [DONE]\n\n",
7401+
"status": 200,
7402+
"statusText": "OK",
7403+
"headers": {
7404+
"access-control-expose-headers": "X-Debug-Trace-ID",
7405+
"alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000",
7406+
"cache-control": "no-cache, no-store, no-transform, must-revalidate, private, max-age=0",
7407+
"content-type": "text/event-stream",
7408+
"expires": "Thu, 01 Jan 1970 00:00:00 UTC",
7409+
"pragma": "no-cache",
7410+
"server": "envoy",
7411+
"transfer-encoding": "chunked",
7412+
"vary": "Origin"
7413+
}
7414+
}
7415+
},
7416+
"545bf4e8393bc07dedb7c66d13846ff8264a49e909117c3c93ae35e30e705cbb": {
7417+
"url": "https://api.cohere.com/v2/chat",
7418+
"init": {
7419+
"headers": {
7420+
"Content-Type": "application/json"
7421+
},
7422+
"method": "POST",
7423+
"body": "{\"messages\":[{\"role\":\"user\",\"content\":\"Complete this sentence with words, one plus one is equal \"}],\"model\":\"command-r7b-12-2024\"}"
7424+
},
7425+
"response": {
7426+
"body": "{\"id\":\"cd9a6a0f-5e4c-411f-9604-fd4bdffc3052\",\"message\":{\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"One plus one is equal to two.\"}]},\"finish_reason\":\"COMPLETE\",\"usage\":{\"billed_units\":{\"input_tokens\":11,\"output_tokens\":8},\"tokens\":{\"input_tokens\":507,\"output_tokens\":10}}}",
7427+
"status": 200,
7428+
"statusText": "OK",
7429+
"headers": {
7430+
"access-control-expose-headers": "X-Debug-Trace-ID",
7431+
"alt-svc": "h3=\":443\"; ma=2592000,h3-29=\":443\"; ma=2592000",
7432+
"cache-control": "no-cache, no-store, no-transform, must-revalidate, private, max-age=0",
7433+
"content-type": "application/json",
7434+
"expires": "Thu, 01 Jan 1970 00:00:00 UTC",
7435+
"num_chars": "2635",
7436+
"num_tokens": "19",
7437+
"pragma": "no-cache",
7438+
"server": "envoy",
7439+
"vary": "Origin"
7440+
}
7441+
}
73897442
}
73907443
}

0 commit comments

Comments
 (0)