from fastapi import FastAPI from transformers import AutoModelForCausalLM, AutoTokenizer import torch app = FastAPI() model_name = "unsloth/Qwen2.5-7B-bnb-4bit" # Change to your model model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16) tokenizer = AutoTokenizer.from_pretrained(model_name) @app.get("/") def read_root(): return {"message": "Chat API is running!"} @app.post("/chat") def chat(prompt: str): inputs = tokenizer(prompt, return_tensors="pt") outputs = model.generate(**inputs, max_new_tokens=100) response = tokenizer.decode(outputs[0], skip_special_tokens=True) return {"response": response}