improvements

Dockerfile

@@ -8,6 +8,6 @@ RUN pip install --upgrade pip && \
 COPY .model .model
 COPY serve.py .
 
-EXPOSE 8000
+EXPOSE 8482
 
 CMD ["python", "serve.py"]

README.md

@@ -6,7 +6,7 @@ Tool to assist developers when writing courses.
 
 2. Run the container:
 
-    podman run --rm -ti -p 8482:8000 quay.io/jramcast/text-generator
+    podman run --rm -ti -p 8482:8482 quay.io/jramcast/text-generator
 
 3. Install the extension in VSCode
 

extension/rht-text-generator/README.old.md

@@ -7,7 +7,7 @@
 
 2. Run the container:
 
-    podman run --rm -ti -p 8482:8000 quay.io/jramcast/text-generator
+    podman run --rm -ti -p 8482:8482 quay.io/jramcast/text-generator
 
 3. Package the extension:
 

extension/rht-text-generator/package.json

@@ -44,7 +44,7 @@
 			"properties": {
 				"rht-text-generator.lines": {
 					"type": "number",
-					"default": 3,
+					"default": 5,
 					"description": "Number of lines to pass to the model"
 				},
 				"rht-text-generator.length": {

extension/rht-text-generator/src/extension.ts

@@ -6,6 +6,8 @@ import { exec } from "child_process";
 import { COMPLETION_TRIGGERS } from "./triggers";
 
 
+const leadingSpacesRegex = /^[^\S\r\n]+/;
+
 
 // this method is called when your extension is activated
 // your extension is activated the very first time the command is executed
@@ -41,25 +43,32 @@ async function getCompletionsListItemsFor(
 ): Promise<vscode.CompletionItem[]> {
 
 	const config = vscode.workspace.getConfiguration("rht-text-generator");
-	const MAX_LINES: number = config.get("lines") || 3;
-	const lines = [];
-
-	for (let lineOffset = 0; lineOffset < MAX_LINES; lineOffset++ ) {
-		const lineNumber = Math.max(0, position.line - lineOffset);
-		const line = document.lineAt(lineNumber).text.trimEnd();
-		lines.unshift(line);
-	}
-	const text = lines.join("\n");
+	const text = getText(document, position);
 
 	const predictionLength: number = config.get("length") || 3;
 
 	const server: string = config.get("server") || "";
 	let suggestions: string[] = await generateSuggestions(text, predictionLength, server);
 
-	return suggestions.map(suggestion => {
-		const tail = suggestion.replace(text, "").trim();
-		return new vscode.CompletionItem(tail);
-	});
+	return suggestions.map((suggestion, idx) => {
+		const item = new vscode.CompletionItem(suggestion.replace(leadingSpacesRegex, ""));
+		item.sortText = idx.toString();
+		return item;
+	}).filter(item => item.label.trim().length > 0);
+}
+
+/**
+ * Get the text from the current position back to a specific number of lines
+ */
+function getText(document: vscode.TextDocument, position: vscode.Position): string  {
+	const config = vscode.workspace.getConfiguration("rht-text-generator");
+	const MAX_LINES: number = config.get("lines") || 3;
+	const currentLineNumber = position.line;
+	const rangeEnd = position;
+	const rangeStartLine = Math.max(0, currentLineNumber - MAX_LINES);
+	const rangeStart = new vscode.Position(rangeStartLine, 0);
+	const text = document.getText(new vscode.Range(rangeStart, rangeEnd));
+	return text;
 }
 
 async function generateSuggestions(line: string, predictionLength: number, server: string) {

extension/rht-text-generator/src/triggers.ts

@@ -4,7 +4,6 @@ export const COMPLETION_TRIGGERS = [
     ",",
     ";",
     "-",
-    "\n",
     "(",
     ")",
     "{",

serve.py

@@ -1,4 +1,3 @@
-from pprint import pprint
 from sanic import Sanic
 from sanic.response import json
 from transformers import pipeline, set_seed, GPT2Tokenizer
@@ -18,32 +17,34 @@ async def test(request):
     num_predicted_tokens = int(request.args.get("length", 3))
     no_topp = request.args.get("no_top", False)
 
-    num_tokens = tokenizer(text, return_length=True)["length"]
+    tokens = tokenizer(text, return_length=True)
+    num_tokens = tokens["length"]
     max_length = num_tokens + num_predicted_tokens
 
-    print("Text:", text)
-
     kargs = {
         "do_sample": True,
         "top_k": max_length,
-        "top_p": 0.9
+        "top_p": 0.92
     }
+
     if no_topp:
         kargs = {}
 
     # For info about the args: https://huggingface.co/blog/how-to-generate
     predictions = generator(
         text,
         max_length=max_length,
-        num_return_sequences=4,
+        num_return_sequences=5,
+        output_scores=True,
+        return_full_text=False,
         **kargs
     )
 
     result = [p["generated_text"] for p in predictions]
 
-    pprint(result)
+    print("Predictions:", result)
 
     return json(result)
 
 if __name__ == '__main__':
-    app.run(host="0.0.0.0")
+    app.run(host="0.0.0.0", port=8482)