app.py

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

class Request(BaseModel):
    prompt: str 
    response: str = Field(description="")
    do_sample: bool=True,
    top_k: int =1,
    temperature: float=0.9,
    max_new_tokens: int=500,
    repetition_penalty: float=1.5,

app = FastAPI()

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

@app.post("/generate")
def generate(req: Request):
    model_name_or_id = "AI4Chem/ChemLLM-7B-Chat"

    model = AutoModelForCausalLM.from_pretrained(model_name_or_id,trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_id,trust_remote_code=True)

    inputs = tokenizer(req.prompt, return_tensors="pt")

    generation_config = GenerationConfig(
        do_sample=req.do_sample,
        top_k=req.top_k,
        temperature=req.temperature,
        max_new_tokens=req.max_new_tokens,
        repetition_penalty=req.repetition_penalty,
        pad_token_id=tokenizer.eos_token_id
    )

    outputs = model.generate(**inputs, generation_config=generation_config)
    # print(tokenizer.decode(outputs[0], skip_special_tokens=True))

    return {"text": tokenizer.decode(outputs[0], skip_special_tokens=True)}