from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

app = FastAPI()

model_name = "THUDM/chatglm2-6b"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    device_map="auto",
    torch_dtype=torch.float16
)
model.eval()

class InputPrompt(BaseModel):
    prompt: str

@app.post("/predict")
def predict(data: InputPrompt):
    try:
        inputs = tokenizer(data.prompt, return_tensors="pt").to(model.device)
        outputs = model.generate(**inputs, max_length=512, use_cache=False)
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return {"result": result}
    except Exception as e:
        return {"error": str(e)}