{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 1 - INSTALL DEPENDENCIES"
      ],
      "metadata": {
        "id": "y4paEjBWXqKf"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7DhjeVpHWxPY"
      },
      "outputs": [],
      "source": [
        "!pip install -q transformers peft accelerate bitsandbytes huggingface_hub"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 2 - MOUNT GOOGLE DRIVE\n"
      ],
      "metadata": {
        "id": "vAN1wKnhXtTb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "collapsed": true,
        "id": "avi4pczBXurI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 3 - IMPORT LIBRARIES\n"
      ],
      "metadata": {
        "id": "Uj46967cXwbC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "from peft import PeftModel\n",
        "from huggingface_hub import login"
      ],
      "metadata": {
        "id": "19dmQJ77Xxzo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 4 - HUGGING FACE AUTHENTICATION\n"
      ],
      "metadata": {
        "id": "k50Y8YDLXz3C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import getpass\n",
        "HF_TOKEN = getpass.getpass(\"Enter your Hugging Face token: \")\n",
        "os.environ[\"HF_TOKEN\"] = HF_TOKEN"
      ],
      "metadata": {
        "id": "NfVn6y3bX0Tr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 5 - VISABUDDY CLASS DEFINITION\n"
      ],
      "metadata": {
        "id": "ZWXng3zhX2Re"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# IMPROVE VISA BUDDY WITH SYSTEM PROMPT\n",
        "class ImprovedVisaBuddy:\n",
        "    def __init__(self, model_path=\"/content/drive/MyDrive/visa_buddy_model\"):\n",
        "        self.model_path = model_path\n",
        "        self.model = None\n",
        "        self.tokenizer = None\n",
        "        self.load_model()\n",
        "\n",
        "    def load_model(self):\n",
        "        print(\"🔄 Loading Improved Visa Buddy model...\")\n",
        "\n",
        "        base_model = \"meta-llama/Llama-3.1-8B-Instruct\"\n",
        "\n",
        "        try:\n",
        "            self.tokenizer = AutoTokenizer.from_pretrained(\n",
        "                base_model,\n",
        "                token=os.environ[\"HF_TOKEN\"]\n",
        "            )\n",
        "            if self.tokenizer.pad_token_id is None:\n",
        "                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id\n",
        "\n",
        "            self.model = AutoModelForCausalLM.from_pretrained(\n",
        "                base_model,\n",
        "                torch_dtype=torch.float16,\n",
        "                device_map=\"auto\",\n",
        "                load_in_8bit=True,\n",
        "                token=os.environ[\"HF_TOKEN\"]\n",
        "            )\n",
        "\n",
        "            if os.path.exists(self.model_path):\n",
        "                self.model = PeftModel.from_pretrained(self.model, self.model_path)\n",
        "                print(\"✅ Fine-tuned adapter loaded!\")\n",
        "            else:\n",
        "                print(\"⚠️  Using base model (no fine-tuned adapter found)\")\n",
        "\n",
        "            self.model.eval()\n",
        "            print(\"✅ Improved Visa Buddy model loaded successfully!\")\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f\"❌ Error loading model: {e}\")\n",
        "\n",
        "    def chat(self, message, max_length=500):\n",
        "        if self.model is None:\n",
        "            return \"Model not loaded.\"\n",
        "\n",
        "        # SYSTEM PROMPT - This is the key improvement!\n",
        "        system_prompt = \"\"\"You are Visa Buddy, a specialized AI assistant for Canadian immigration and visa information.\n",
        "\n",
        "YOUR ROLE:\n",
        "- Provide clear, direct, and accurate information about Canadian visas, immigration, and related processes\n",
        "- Focus exclusively on Canadian immigration topics\n",
        "- Be concise and factual - avoid unnecessary fluff or conversational filler\n",
        "- If asked about non-visa topics, politely redirect to Canadian immigration\n",
        "- Use bullet points and clear structure when listing requirements\n",
        "- Always cite specific programs, requirements, and processing times when possible\n",
        "\n",
        "RESPONSE STYLE:\n",
        "- Direct and professional\n",
        "- Structured and easy to read\n",
        "- Factual and up-to-date\n",
        "- Focused on actionable information\n",
        "\n",
        "NEVER:\n",
        "- Engage in general chit-chat\n",
        "- Provide opinions or personal advice\n",
        "- Discuss topics unrelated to Canadian immigration\n",
        "- Use vague or uncertain language\"\"\"\n",
        "\n",
        "        # Format the prompt with system message\n",
        "        prompt = f\"\"\"<|start_header_id|>system<|end_header_id|>\n",
        "\n",
        "{system_prompt}<|eot_id|>\n",
        "<|start_header_id|>user<|end_header_id|>\n",
        "\n",
        "{message}<|eot_id|>\n",
        "<|start_header_id|>assistant<|end_header_id|>\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "        inputs = self.tokenizer(prompt, return_tensors=\"pt\").to(self.model.device)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            outputs = self.model.generate(\n",
        "                **inputs,\n",
        "                max_new_tokens=max_length,\n",
        "                temperature=0.7,\n",
        "                do_sample=True,\n",
        "                pad_token_id=self.tokenizer.eos_token_id,\n",
        "                eos_token_id=self.tokenizer.eos_token_id,\n",
        "                repetition_penalty=1.1,\n",
        "                no_repeat_ngram_size=3\n",
        "            )\n",
        "\n",
        "        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "\n",
        "        # Extract only the assistant's response\n",
        "        if \"assistant\" in response:\n",
        "            response = response.split(\"assistant\")[-1].strip()\n",
        "\n",
        "        return response\n",
        "\n",
        "# Replace the existing visa_buddy with improved version\n",
        "visa_buddy = ImprovedVisaBuddy()\n",
        "visa_buddy_instance = visa_buddy\n",
        "print(\"🎯 Improved Visa Buddy with system prompt is ready!\")"
      ],
      "metadata": {
        "id": "Ez7Pl7B4X2vS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 6 - INITIALIZE MODEL\n"
      ],
      "metadata": {
        "id": "uc49YlpMX5CY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "visa_buddy_instance = visa_buddy\n"
      ],
      "metadata": {
        "id": "6X76qJddX6fP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 7 - INSTALL DEPENDENCIES\n"
      ],
      "metadata": {
        "id": "tTWu0iLTgVKC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q fastapi uvicorn python-multipart\n",
        "!npm install -g localtunnel"
      ],
      "metadata": {
        "id": "tb9_uscmgWXQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 8 - FASTAPI SERVER\n"
      ],
      "metadata": {
        "id": "HqqiKrYjgYBV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from fastapi import FastAPI, HTTPException\n",
        "from fastapi.middleware.cors import CORSMiddleware\n",
        "from pydantic import BaseModel\n",
        "import uvicorn\n",
        "import threading\n",
        "\n",
        "class ChatRequest(BaseModel):\n",
        "    message: str\n",
        "    max_length: int = 700\n",
        "\n",
        "class ChatResponse(BaseModel):\n",
        "    response: str\n",
        "\n",
        "app = FastAPI(title=\"Visa Buddy API\", version=\"1.0.0\")\n",
        "\n",
        "app.add_middleware(\n",
        "    CORSMiddleware,\n",
        "    allow_origins=[\"*\"],\n",
        "    allow_credentials=True,\n",
        "    allow_methods=[\"*\"],\n",
        "    allow_headers=[\"*\"],\n",
        ")\n",
        "\n",
        "visa_buddy_instance = None\n",
        "\n",
        "@app.get(\"/\")\n",
        "async def root():\n",
        "    return {\"message\": \"Visa Buddy API is running!\"}\n",
        "\n",
        "@app.post(\"/chat\", response_model=ChatResponse)\n",
        "async def chat_endpoint(request: ChatRequest):\n",
        "    try:\n",
        "        if visa_buddy_instance is None:\n",
        "            raise HTTPException(status_code=500, detail=\"Model not loaded\")\n",
        "\n",
        "        response = visa_buddy_instance.chat(request.message, request.max_length)\n",
        "        return ChatResponse(response=response)\n",
        "    except Exception as e:\n",
        "        raise HTTPException(status_code=500, detail=f\"Error: {str(e)}\")\n",
        "\n",
        "@app.get(\"/health\")\n",
        "async def health_check():\n",
        "    return {\"status\": \"healthy\", \"model_loaded\": visa_buddy_instance is not None}\n",
        "\n",
        "def start_server():\n",
        "    uvicorn.run(app, host=\"0.0.0.0\", port=8000)\n",
        "\n",
        "visa_buddy_instance = visa_buddy\n",
        "print(\"✅ Visa Buddy model ready for API!\")"
      ],
      "metadata": {
        "id": "RNkzVQgZgYy4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# CELL 9 - START SERVER & GET PUBLIC URL\n"
      ],
      "metadata": {
        "id": "d3oQsm6JgZ0K"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import threading\n",
        "import time\n",
        "\n",
        "# Start FastAPI server in background\n",
        "server_thread = threading.Thread(target=start_server, daemon=True)\n",
        "server_thread.start()\n",
        "print(\"🔄 Starting FastAPI server...\")\n",
        "time.sleep(3)\n",
        "\n",
        "print(\"🚀 Run this command in a NEW cell to get your public URL:\")\n",
        "print(\"!lt --port 8000 --subdomain visa-buddy-$(date +%s)\")\n",
        "print(\"\\nOr use this simpler version:\")\n",
        "print(\"!lt --port 8000\")\n",
        "print(\"\\n📝 Copy the URL it provides and use it in your React app!\")"
      ],
      "metadata": {
        "id": "mG1SqJ5_gbVO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64\n",
        "!chmod +x cloudflared-linux-amd64\n",
        "!./cloudflared-linux-amd64 tunnel --url http://localhost:8000"
      ],
      "metadata": {
        "id": "x4eVTE6NlvQH"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}