mjpsm/progress-generation-model-api

1from fastapi import FastAPI2from pydantic import BaseModel3from transformers import AutoTokenizer, AutoModelForCausalLM4import torch5 6app = FastAPI()7 8MODEL_NAME = "mjpsm/progress-generation-model"9 10device = "cuda" if torch.cuda.is_available() else "cpu"11 12tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)13model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)14 15tokenizer.pad_token = tokenizer.eos_token16 17class Request(BaseModel):18    text: str19 20def generate_response(user_input):21    prompt = f"""<|system|>22You describe what progress was achieved in one sentence.23<|user|>24{user_input}25<|assistant|>26"""27 28    inputs = tokenizer(prompt, return_tensors="pt").to(device)29 30    with torch.no_grad():31        outputs = model.generate(32            **inputs,33            max_new_tokens=50,34            temperature=0.6,35            top_p=0.9,36            repetition_penalty=1.2,37            pad_token_id=tokenizer.eos_token_id38        )39 40    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)41    return decoded.split("<|assistant|>")[-1].strip()42 43@app.get("/")44def root():45    return {"message": "Progress Model API running"}46 47@app.post("/predict")48def predict(req: Request):49    result = generate_response(req.text)50    return {"output": result}