Switch to CPU-compatible model

app.py

@@ -1,25 +1,17 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from typing import List
 
 app = FastAPI()
 
-model_id = "unsloth/gemma-2-2b-it-bnb-4bit"
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4"
-)
+model_id = "google/gemma-2b-it"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    quantization_config=bnb_config,
-    device_map="auto"
+    device_map="cpu"
 )
 
 class InferenceRequest(BaseModel):
@@ -38,7 +30,7 @@ def read_root():
 
 @app.post("/infer")
 def infer(request: InferenceRequest):
-    inputs = tokenizer(request.text, return_tensors="pt").to("cuda")
+    inputs = tokenizer(request.text, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=50)
     decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"generated_text": decoded_output}
@@ -47,9 +39,9 @@ def infer(request: InferenceRequest):
 def chat(request: ChatRequest):
     chat_history = [message.dict() for message in request.messages]
     prompt = tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_new_tokens=150)
     decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
     # The output from the model includes the prompt, so we need to remove it.
     response = decoded_output.split("<end_of_turn>\n")[-1]
-    return {"generated_text": response}
+    return {"generated_text": response}

requirements.txt

@@ -2,8 +2,7 @@ fastapi
 uvicorn
 torch
 transformers
-bitsandbytes
 accelerate
 sentencepiece
 python-dotenv
-pydantic
+pydantic