app.py

from fastapi import FastAPI
import uvicorn
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_name = "TheBloke/Wizard-Vicuna-13B-Uncensored-HF"

# Configure 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,               # Enable 4-bit quantization
    bnb_4bit_quant_type="nf4",       # Use 4-bit NormalFloat (optimal)
    bnb_4bit_compute_dtype="float16", # Faster computation with float16
    bnb_4bit_use_double_quant=True   # Extra compression
)

# Load model with quantization
model = AutoModelForCausalLM.from_pretrained(
    model_name,  # Example model
    quantization_config=bnb_config,
    device_map="auto",            # Auto-distribute across GPU/CPU
    trust_remote_code=True        # Required for some models
)

# load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

app = FastAPI()

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

@app.get("/message")
async def message(input: str):
    inputs = tokenizer(input, return_tensors="pt", padding=True, truncation=True)
    
    output = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"], 
        max_new_tokens=100,
    )
    
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    return response

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)