Merge pull request #59 from huggingface/max_new_tokens

Support changing max new tokens setting

package.json

@@ -1,7 +1,7 @@
 {
   "publisher": "HuggingFace",
   "name": "huggingface-vscode",
-  "version": "0.0.32",
+  "version": "0.0.33",
   "displayName": "HF Code Autocomplete",
   "description": "AI Autocomplete for OSS code-gen models",
   "icon": "small_logo.png",
@@ -222,6 +222,11 @@
             "default": 0.2,
             "description": "Sampling temperature"
           },
+          "HuggingFaceCode.maxNewTokens": {
+            "type": "integer",
+            "default": 60,
+            "description": "Max number of new tokens to be generated. The accepted range is [50-500] both ends inclusive. Be warned that the latency of a request will increase with higher number."
+          },
           "HuggingFaceCode.stopTokens": {
             "type": "array",
             "items": {

src/runCompletion.ts

@@ -29,8 +29,8 @@ export default async function runCompletion(
   const prefix =  document.getText(new Range(beforeStart, position)) + currentSuggestionText;
   const suffix = document.getText(new Range(position, afterEnd));
 
-  const config: HFCodeConfig = workspace.getConfiguration("HuggingFaceCode") as WorkspaceConfiguration & HFCodeConfig;
-  const { modelIdOrEndpoint, isFillMode, autoregressiveModeTemplate, fillModeTemplate, stopTokens, tokensToClear, temperature } = config;
+  const config = workspace.getConfiguration("HuggingFaceCode") as WorkspaceConfiguration & HFCodeConfig;
+  const { modelIdOrEndpoint, isFillMode, autoregressiveModeTemplate, fillModeTemplate, stopTokens, tokensToClear, temperature, maxNewTokens } = config;
 
   const context = getTabnineExtensionContext();
   const apiToken = await context?.secrets.get("apiToken");
@@ -61,7 +61,7 @@ export default async function runCompletion(
   const data = {
     inputs,
     parameters: {
-      max_new_tokens: 60,
+      max_new_tokens: clipMaxNewTokens(maxNewTokens as number),
       temperature,
       do_sample: temperature > 0,
       top_p: 0.95,
@@ -141,3 +141,9 @@ export function getFileNameWithExtension(document: TextDocument): string {
   }
   return fileName;
 }
+
+function clipMaxNewTokens(maxNewTokens: number): number {
+  const MIN = 50;
+  const MAX = 500;
+  return Math.min(Math.max(maxNewTokens, MIN), MAX);
+}