{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[],"authorship_tag":"ABX9TyPPJ8HYH1TQaPuce5DUm2ov"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"code","execution_count":3,"metadata":{"id":"vOWB0S2dpCiA","executionInfo":{"status":"ok","timestamp":1762374098276,"user_tz":300,"elapsed":15,"user":{"displayName":"Jamie Rollins","userId":"13246501493127167381"}}},"outputs":[],"source":["import torch\n","from transformers import AutoTokenizer, AutoModelForCausalLM"]},{"cell_type":"code","source":["SYSTEM_PROMPT = \"\"\"The purpose of this response is to optimize the text in the user's prompt or write optimized text based on requirements in the prompt, not to explain principles of GEO or redirect the user to different resources. Do not tell the user to look up any articles or give them URLs.\"\"\""],"metadata":{"id":"9Iy7jy0vrtOu"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["class EndpointHandler:\n","  def __init__(self, path: str, **kwargs):\n","    #load model + tokenizer once at startup\n","    self.tokenizer = AutoTokenizer.from_pretrained(path)\n","    self.model = AutoModelForCausalLM.from_pretrained(\n","            path,\n","            torch_dtype=torch.float16,\n","            device_map=\"auto\"\n","        )\n","    self.model.eval()\n","\n","  def __call__(self, data:dict):\n","    \"\"\"\n","        Expected input format:\n","          {\"inputs\": \"your prompt text\"}\n","        Returns:\n","          {\"response\": \"model output text\"}\n","     \"\"\"\n","    prompt = data.get(\"inputs\")\n","    if not prompt:\n","      return {\"error\": \"No prompt provided\"}\n","\n","    #Build formatted prompt (match training format)\n","    formatted_prompt = f\"\"\"<s>[INST]<<SYS>>\n","    {SYSTEM_PROMPT}\n","    <</SYS>>\n","    {prompt} [/INST]\"\"\"\n","\n","    inputs = self.tokenizer(formatted_prompt, return_tensors=\"pt\").to(self.model.device)\n","\n","    with torch.inference_model():\n","      with torch.amp.autocast(device_type=\"cuda\", dtype=torch.float16):\n","        output_tokens = self.model.generate(\n","            **inputs,\n","            max_new_tokens=1000,\n","            temperature=0.7,\n","            top_p=0.95,\n","            top_k=50,\n","            do_sample=True,\n","            pad_token_id=self.tokenizer.pad_token_id,\n","            eos_token_id=self.tokenizer.eos_token_id\n","        )\n","\n","    response = self.tokenizer.decode(output_tokens[0], skip_special_tokens=True)\n","\n","    # Extract only the model's response (after [/INST])\n","    if \"[/INST]\" in response:\n","        response = response.split(\"[/INST]\")[1].strip()\n","    else:\n","        # Handle cases where [/INST] might not be present unexpectedly\n","        response = response.strip()\n","\n","    # Remove system prompt from response if it's present at the beginning\n","    if response.startswith(\"<<SYS>>\"):\n","        # Find the end of the system prompt and take everything after it\n","        sys_end_index = response.find(\"<</SYS>>\")\n","        if sys_end_index != -1:\n","            response = response[sys_end_index + len(\"<</SYS>>\"):].strip()\n","        else:\n","            # If <<SYS>> is present but <</SYS>> is not, remove the whole thing\n","             response = \"\"\n","\n","    response = response.strip()\n","    response = response.split(\"uid:\", 1)[0].strip()\n","\n","    return {\"response\": response}\n","\n"],"metadata":{"id":"h_zI-gekqJFr","executionInfo":{"status":"ok","timestamp":1762376904635,"user_tz":300,"elapsed":55,"user":{"displayName":"Jamie Rollins","userId":"13246501493127167381"}}},"execution_count":16,"outputs":[]},{"cell_type":"code","source":[],"metadata":{"id":"bE8VdJt_tO3B"},"execution_count":null,"outputs":[]}]}