Update lambda to use new model

Author: Nishant2022

src/backend/Dockerfile

@@ -7,7 +7,10 @@ COPY lambda_requirements.txt ${LAMBDA_TASK_ROOT}
 RUN pip install --compile --no-cache-dir -r lambda_requirements.txt
 
 # Copy model
-COPY trained/gpt2-728 ${LAMBDA_TASK_ROOT}/trained/gpt2-728
+COPY trained/mini-copilot ${LAMBDA_TASK_ROOT}/trained/mini-copilot
+
+# Copy tokenizer
+COPY trained/mini-copilot-tokenizer ${LAMBDA_TASK_ROOT}/trained/mini-copilot-tokenizer
 
 # Copy function code
 COPY lambda.py ${LAMBDA_TASK_ROOT}

src/backend/lambda.py

@@ -4,23 +4,22 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-model = AutoModelForCausalLM.from_pretrained("trained/gpt2-728")
+tokenizer = AutoTokenizer.from_pretrained(
+    "trained/mini-copilot-tokenizer/tokenizer_10M")
+model = AutoModelForCausalLM.from_pretrained("trained/mini-copilot/gpt2-large")
 model.to(device)
 model.eval()
 
+pipe = pipeline("text-generation", model=model,
+                tokenizer=tokenizer, device=device, max_new_tokens=32)
 
-# Function to predict the next token
-def predict_next_token(input_text, model=model, tokenizer=tokenizer, max_length=50):
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
-
-    predicted_text = pipe(input_text)[0]["generated_text"]
 
-    input_text_len = len(input_text)
-
-    return predicted_text[input_text_len:]
+# Function to predict the next token
+def get_completion(inp: str) -> str:
+    return pipe(inp)[0]["generated_text"][len(inp):]
 
 
+# Lambda handler
 def handler(event, context):
     try:
         if event.get('isBase64Encoded', False):
@@ -29,4 +28,4 @@ def handler(event, context):
             body = event['body']
     except (KeyError, json.JSONDecodeError) as e:
         return {"statusCode": 400, "body": f"Error processing request: {str(e)}"}
-    return {"statusCode": 200, "body": predict_next_token(body)}
+    return {"statusCode": 200, "body": get_completion(body)}