Skip to content

Commit

Permalink
Merge branch 'main' into auto-completion
Browse files Browse the repository at this point in the history
  • Loading branch information
@Wzixiao
Wzixiao committed Oct 31, 2023
2 parents dc40985 + b5a0ada commit 8f2390c
Show file tree
Hide file tree
Showing 4 changed files with 198 additions and 116 deletions.
252 changes: 149 additions & 103 deletions packages/jupyter-ai/src/bigcode-Inline-completion-provider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,20 +11,11 @@ import { NotebookPanel } from '@jupyterlab/notebook';
import { nullTranslator, TranslationBundle } from '@jupyterlab/translation';
import { retrieveNotebookContentUntilCursor } from './utils/cell-context';

import bigcodeRequestInstance from './utils/bigcode-request';
import bigcodeRequestInstance, {
BigCodeServiceStreamResponseItem
} from './utils/bigcode-request';
import CodeCompletionContextStore from './contexts/code-completion-context-store';

// type BigCodeStream = {
// token: {
// id: number;
// text: string;
// logprob: number;
// special: boolean;
// };
// generated_text: string | null;
// details: null;
// };

export class BigcodeInlineCompletionProvider
implements IInlineCompletionProvider
{
Expand Down Expand Up @@ -105,7 +96,7 @@ export class BigcodeInlineCompletionProvider
return '';
}

async shortCutHandler(
async shortCutCompletionHandler(
request: CompletionHandler.IRequest,
context: IInlineCompletionContext
): Promise<IInlineCompletionList<IInlineCompletionItem>> {
Expand Down Expand Up @@ -169,7 +160,6 @@ export class BigcodeInlineCompletionProvider
}

const result = await this.simulateSingleRequest(prompt);
console.debug(result);
if (result === '<debounce>') {
return { items: [] };
}
Expand Down Expand Up @@ -205,7 +195,10 @@ export class BigcodeInlineCompletionProvider
return { items: [] };
}

if (!CodeCompletionContextStore.accessToken) {
if (
!CodeCompletionContextStore.accessToken &&
!CodeCompletionContextStore.enableMockTest
) {
alert('Huggingface Access Token not set.');
return { items: [] };
}
Expand All @@ -215,7 +208,7 @@ export class BigcodeInlineCompletionProvider
}

if (context.triggerKind === 0) {
return await this.shortCutHandler(request, context);
return await this.shortCutCompletionHandler(request, context);
}

return { items: [] };
Expand Down Expand Up @@ -259,7 +252,6 @@ export class BigcodeInlineCompletionProvider
if (this._callCounter === currentCallCount && !this._requesting) {
this._callCounter = 0;
this._requesting = true;
// resolve('"""This is the result of a simulated automatic request"""');
resolve('<auto_stream>');
this._timeoutId = null;
} else {
Expand All @@ -271,114 +263,168 @@ export class BigcodeInlineCompletionProvider

async *stream(
token: string
): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
if (!CodeCompletionContextStore.enableMockTest) {
yield* this.completionStream(token);
} else {
yield* this.mockCompletionStream(token);
}
}

private async *mockCompletionStream(
token: string
): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
this._requesting = true;

if (this._requestMode === 0) {
const testResultText =
'_world():\n print("Hello World!")\nhello_world()';
for (let i = 1; i <= testResultText.length; i++) {
await this.delay(25);
yield* this.mockKeypressCompletionStream();
} else if (this._requestMode === 1) {
yield* this.mockAutomaticCompletionStream(token);
}

if (this._streamStop) {
console.debug('_streamStop');
this.setRequestFinish(false);
this.setRequestFinish(false);
}

yield {
response: {
isIncomplete: false,
insertText: this._lastRequestInfo.insertText
}
};
return;
}
private async *mockKeypressCompletionStream(): AsyncGenerator<
{ response: IInlineCompletionItem },
undefined,
unknown
> {
const testResultText =
'"""This is the first line generated by the mockKeypressCompletionStreamr"""\n """This is the second line generated by the mockKeypressCompletionStream"""';

const insertChar = testResultText.slice(i - 1, i);
this._lastRequestInfo.insertText += insertChar;
for (let i = 1; i <= testResultText.length; i++) {
await this.delay(10);

if (this._streamStop) {
console.debug('_streamStop');
yield {
response: {
isIncomplete: i !== testResultText.length - 1,
isIncomplete: false,
insertText: this._lastRequestInfo.insertText
}
};
return;
}
this.setRequestFinish(false);

const insertChar = testResultText.slice(i - 1, i);
this._lastRequestInfo.insertText += insertChar;

yield {
response: {
isIncomplete: i !== testResultText.length - 1,
insertText: this._lastRequestInfo.insertText
}
};
}
}

private async *mockAutomaticCompletionStream(
token: string
): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
await this.delay(1000);
const insertText =
'"""This is the first line that automatically triggers the mockAutomaticCompletionStream"""\n """This is the second line that automatically triggers the mockAutomaticCompletionStream"""';
this._lastRequestInfo.insertText = insertText;

yield {
response: {
token,
isIncomplete: false,
insertText: insertText
}
};
}

private async *completionStream(
token: string
): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
this._requesting = true;

if (this._requestMode === 0) {
yield* this.keypressCompletionStream(token);
} else if (this._requestMode === 1) {
await this.delay(1000);
const insertText =
'"""This is the result of a simulated automatic request"""';
this._lastRequestInfo.insertText = insertText;
yield* this.automaticCompletionStream(token);
}

this.setRequestFinish(false);
}

async *keypressCompletionStream(
token: string
): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
if (token === '') {
return;
}

this._requesting = true;
this._finish = false;
let reponseData: ReadableStream<Uint8Array> | null = null;

try {
reponseData = await bigcodeRequestInstance.fetchStream(true);
} catch {
yield {
response: {
token,
isIncomplete: false,
insertText: insertText
insertText: ''
}
};
this.setRequestFinish(false);
this.setRequestFinish(true);
return;
}

const decoder = new TextDecoder();
const reader = reponseData.getReader();

while (true) {
const { value, done } = await reader.read();

if (done || this._streamStop) {
return;
}

const strValue = decoder.decode(value, { stream: true });

const jsonStrList = strValue.split('data:');

for (const chunk of jsonStrList) {
if (chunk !== '') {
const chunkData = JSON.parse(
chunk
) as BigCodeServiceStreamResponseItem;
const done = chunkData.token.special;

if (!done) {
this._lastRequestInfo.insertText += chunkData.token.text;
}

yield {
response: {
isIncomplete: !done,
insertText: this._lastRequestInfo.insertText
}
};
}
}
}
}

// async *stream(
// token: string
// ): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
// if (token === '' || this._requesting) {
// return;
// }

// this._requesting = true;
// this._finish = false;
// let reponseData: ReadableStream<Uint8Array> | null = null;

// try {
// reponseData = await BigcodeInstance.fetchStream();
// } catch {
// yield {
// response: {
// isIncomplete: false,
// insertText: ''
// }
// };
// this.setRequestFinish(true);
// return;
// }

// const decoder = new TextDecoder();
// const reader = reponseData.getReader();

// while (true) {
// const { value, done } = await reader.read();

// if (done || this._streamStop) {
// this.setRequestFinish(false);
// break;
// }

// const strValue = decoder.decode(value, { stream: true });

// const jsonStrList = strValue.split('data:');

// for (const chunk of jsonStrList) {
// if (chunk !== '') {
// const chunkData = JSON.parse(chunk) as BigCodeStream;
// const done = chunkData.token.special;

// if (done) {
// this.setRequestFinish(false);
// } else {
// this._lastRequestInfo.insertText += chunkData.token.text;
// }

// yield {
// response: {
// isIncomplete: !chunkData.token.special,
// insertText: this._lastRequestInfo.insertText
// }
// };
// }
// }
// }
// }
async *automaticCompletionStream(
token: string
): AsyncGenerator<{ response: IInlineCompletionItem }, undefined, unknown> {
try {
const reponseData = await bigcodeRequestInstance.fetchStream(false);

yield {
response: {
token,
isIncomplete: false,
insertText: reponseData[0].generated_text
}
};
} catch {
return;
}
}
}
17 changes: 16 additions & 1 deletion packages/jupyter-ai/src/components/bigcode-settings.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,10 @@ export const BigCodeSetting = observer(() => {
CodeCompletionContextstore.toggleCodeCompletion();
};

const toggleMockTestWrapper = () => {
CodeCompletionContextstore.toggleCMockTest();
};

useEffect(() => {
CodeCompletionContextstore.setBigcodeUrl(
'https://api-inference.huggingface.co/models/bigcode/starcoderbase/'
Expand Down Expand Up @@ -75,9 +79,20 @@ export const BigCodeSetting = observer(() => {
onChange={toggleCodeCompletionWrapper}
/>
</h1>
{CodeCompletionContextstore.enableCodeCompletion && (
<h1 className="jp-ai-ChatSettings-header bigcode-setting-level-1-title">
Enable mock test code completion
<Switch
checked={CodeCompletionContextstore.enableMockTest}
onChange={toggleMockTestWrapper}
/>
</h1>
)}

<div
className={
CodeCompletionContextstore.enableCodeCompletion
CodeCompletionContextstore.enableCodeCompletion &&
!CodeCompletionContextstore.enableMockTest
? 'code-completion-setting-entered'
: 'code-completion-setting-exiting'
}
Expand Down
Loading

0 comments on commit 8f2390c

Please sign in to comment.