Add exception handling for incompatible models

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

main.py

@@ -1,4 +1,5 @@
 import zipfile
+from abc import ABC
 from typing import Optional, Union
 from awq import AutoAWQForCausalLM
 from pydantic import BaseModel, Field
@@ -17,37 +18,61 @@ app = FastAPI(title="Huggingface Safetensor Model Converter to AWQ", version="0.
 ### -------
 
 ### DTO Definitions
-class QuantizationConfig(BaseModel):
-    zero_point: Optional[bool] = Field(True, description="Use zero point quantization")
-    q_group_size: Optional[int] = Field(128, description="Quantization group size")
-    w_bit: Optional[int] = Field(4, description="Weight bit")
-    version: Optional[str] = Field("GEMM", description="Quantization version")
-
-class ConvertRequest(BaseModel):
+class QuantizationConfig(ABC, BaseModel):
+    pass
+class ConvertRequest(ABC, BaseModel):
     hf_model_name: str
     hf_tokenizer_name: Optional[str] = Field(None, description="Hugging Face tokenizer name. Defaults to hf_model_name")
     hf_token: Optional[str] = Field(None, description="Hugging Face token for private models")
     hf_push_repo: Optional[str] = Field(None, description="Hugging Face repo to push the converted model. If not provided, the model will be downloaded only.")
-    quantization_config: QuantizationConfig = Field(QuantizationConfig(), description="Quantization configuration")
 ### -------
 
+### Quantization Configurations
+class AWQQuantizationConfig(QuantizationConfig):
+    zero_point: Optional[bool] = Field(True, description="Use zero point quantization")
+    q_group_size: Optional[int] = Field(128, description="Quantization group size")
+    w_bit: Optional[int] = Field(4, description="Weight bit")
+    version: Optional[str] = Field("GEMM", description="Quantization version")
+
+class GPTQQuantizationConfig(QuantizationConfig):
+    pass
+
+class GGUFQuantizationConfig(QuantizationConfig):
+    pass
+class AWQConvertionRequest(ConvertRequest):
+    quantization_config: Optional[AWQQuantizationConfig] = Field(
+        default_factory=lambda: AWQQuantizationConfig(),
+        description="AWQ quantization configuration"
+    )
+
+class GPTQConvertionRequest(ConvertRequest):
+    quantization_config: Optional[GPTQQuantizationConfig] = Field(
+        default_factory=lambda: GPTQQuantizationConfig(),
+        description="GPTQ quantization configuration"
+    )
+
+class GGUFConvertionRequest(ConvertRequest):
+    quantization_config: Optional[GGUFQuantizationConfig] = Field(
+        default_factory=lambda: GGUFQuantizationConfig(),
+        description="GGUF quantization configuration"
+    )
+### -------
 
 @app.get("/", include_in_schema=False)
 def redirect_to_docs():
     return RedirectResponse(url='/docs')
 
 ### FastAPI Endpoints
-@app.get("/health")
-def read_root():
-    return {"status": "ok"}
-
-@app.post("/convert", response_model=None)
-def convert(request: ConvertRequest)->Union[FileResponse, dict]:
+@app.post("/convert_awq", response_model=None)
+def convert(request: AWQConvertionRequest)->Union[FileResponse, dict]:
     model = AutoAWQForCausalLM.from_pretrained(request.hf_model_name)
     tokenizer = AutoTokenizer.from_pretrained(request.hf_tokenizer_name or request.hf_model_name, trust_remote_code=True)
-
-    model.quantize(tokenizer, quant_config=request.quantization_config.model_dump())
-
+    
+    try:
+        model.quantize(tokenizer, quant_config=request.quantization_config.model_dump())
+    except TypeError as e:
+        raise HTTPException(status_code=400, detail=f"Is this model supported by AWQ Quantization? Check:https://github.com/mit-han-lab/llm-awq?tab=readme-ov-file {e}")
+    
     if request.hf_push_repo:
         model.save_quantized(request.hf_push_repo)
         tokenizer.save_pretrained(request.hf_push_repo)
@@ -72,4 +97,16 @@ def convert(request: ConvertRequest)->Union[FileResponse, dict]:
             )
         
     
-    raise HTTPException(status_code=500, detail="Failed to convert model")
+    raise HTTPException(status_code=500, detail="Failed to convert model")
+
+@app.post("/convert_gpt_q", response_model=None)
+def convert_gpt_q(request: ConvertRequest)->Union[FileResponse, dict]:
+    raise HTTPException(status_code=501, detail="Not implemented yet")
+
+@app.post("/convert_gguf", response_model=None)
+def convert_gguf(request: ConvertRequest)->Union[FileResponse, dict]:
+    raise HTTPException(status_code=501, detail="Not implemented yet")
+
+@app.get("/health")
+def read_root():
+    return {"status": "ok"}

requirements.txt

@@ -9,4 +9,7 @@ fastapi[standard]
 transformers
 huggingface_hub
 autoawq[kernels]
-starlette>=0.46.2
+starlette>=0.46.2
+pytest
+requests
+environs

tests/.env.example

@@ -0,0 +1,2 @@
+ENDPOINT=
+HF_TOKEN=

tests/test_convertion.py

@@ -0,0 +1,31 @@
+import pytest
+import requests
+from environs import Env
+from huggingface_hub import login
+
+env = Env()
+env.read_env(override=True)
+
+@pytest.mark.parametrize("model_name", [
+    "gpt2",
+])
+def test_convert_download(model_name):
+    if env.str("HF_TOKEN"):
+        login(token=env("HF_TOKEN"))
+    
+    response = requests.post(
+        env.str("ENDPOINT"),
+        json={
+            "hf_model_name": model_name,
+            "hf_tokenizer_name": model_name,
+            "hf_push_repo": None,       
+        }
+    )
+
+    response.raise_for_status()
+
+    assert response.content_type == 'application/zip'
+
+
+def test_convert_push():
+    pass