Update app.py

app.py

@@ -7,7 +7,7 @@ import time
 from fastapi import FastAPI, HTTPException, Depends
 from fastapi.security.api_key import APIKeyHeader
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+from transformers import AutoModelForCausalLM, AutoConfig, GPT2TokenizerFast
 from starlette.status import HTTP_403_FORBIDDEN, HTTP_503_SERVICE_UNAVAILABLE
 
 # --- 1. GLOBAL VARIABLES ---
@@ -16,17 +16,16 @@ model = None
 generated_keys = {}
 
 # --- 2. CONFIGURATION ---
-MODEL_PATH = "/app/model"
+MODEL_PATH = "/app/model"  # folder with vocab.json, merges.txt, model weights
 API_KEY_NAME = "X-API-Key"
 api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
-app = FastAPI(title="Overflow-111.7B Self-Registering API")
+app = FastAPI(title="Overflow-111.7B API")
 
 # --- 3. MODEL & TOKENIZER LOADING ---
 print("Starting Engine: Initializing Self-Registration...")
 
 try:
-    # Add model path to sys.path
     if MODEL_PATH not in sys.path:
         sys.path.insert(0, MODEL_PATH)
 
@@ -34,32 +33,25 @@ try:
     import configuration_overflow
     import modeling_overflow
 
-    # Load config from configuration_overflow
+    # Register config and model
     conf_class = configuration_overflow.OverflowConfig
+    model_class = modeling_overflow.OverflowForCausalLM
+
     AutoConfig.register("overflow", conf_class)
-    print("Successfully registered 'overflow' config.")
-
-    # Find the correct model class that uses conf_class
-    model_class_candidates = [
-        getattr(modeling_overflow, name)
-        for name in dir(modeling_overflow)
-        if hasattr(getattr(modeling_overflow, name), "config_class") 
-        and getattr(modeling_overflow, name).config_class == conf_class
-    ]
-    if not model_class_candidates:
-        raise Exception("No model class matches the OverflowConfig class.")
-    model_class = model_class_candidates[0]
-
-    # Register model class
     AutoModelForCausalLM.register(conf_class, model_class)
     print(f"Successfully registered model class {model_class.__name__}.")
 
-    # Load tokenizer
-    print("Loading Tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
+    # Load GPT2-style tokenizer from vocab.json + merges.txt
+    tokenizer = GPT2TokenizerFast(
+        vocab_file=f"{MODEL_PATH}/vocab.json",
+        merges_file=f"{MODEL_PATH}/merges.txt",
+        unk_token="",
+        bos_token="",
+        eos_token=""
+    )
+    print("Tokenizer loaded successfully.")
 
-    # Load model
-    print("Loading Model Weights...")
+    # Load model weights
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_PATH,
         trust_remote_code=True,
@@ -67,7 +59,7 @@ try:
         torch_dtype=torch.bfloat16,
         low_cpu_mem_usage=True
     )
-    print("Engine Status: ONLINE")
+    print("Model loaded successfully. Engine Status: ONLINE")
 
 except Exception as e:
     print(f"CRITICAL LOADING ERROR: {e}")