ApolloPro7/Generate_Script

1import os2from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration3from peft import PeftModel4from fastapi import FastAPI5from pydantic import BaseModel6from huggingface_hub import login7 8login(token=os.getenv("HF_TOKEN"))9print("Hugging Face Successfully Login!")10 11app = FastAPI()12 13# Load fine-tuned model and tokenizer14# tokenizer = AutoTokenizer.from_pretrained("./llama2-7b", local_files_only=True)15# base_model = AutoModelForCausalLM.from_pretrained("./llama2-7b", local_files_only=True)16# model = PeftModel.from_pretrained(base_model, "./checkpoint-5400", local_files_only=True)17 18 19tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")20model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")21 22# Define data structure of parameters23class PromptInput(BaseModel):24    prompt: str25 26# define API interface27# @app.post("/generate")28# def generate_script(input: PromptInput):29#     print("Starts Generating!")30#     inputs = tokenizer(input.prompt, return_tensors="pt")31#     print("Inputs Tokenized! Generating Begins~")32#     outputs = model.generate(**inputs, max_new_tokens=200)33#     print("Generating Succeed!")34#     result = tokenizer.decode(outputs[0], skip_special_tokens=True)35#     print("Results formed!")36#     return {"generated_script": result}37 38@app.post("/generate")39def generate_script(input: PromptInput):40    print("Starts Generating!")41    inputs = tokenizer(input.prompt, return_tensors="pt").input_ids42    print("Inputs Tokenized! Generating Begins~")43    outputs = model.generate(inputs)44    print("Generating Succeed!")45    result = tokenizer.decode(outputs[0])46    print("Results formed!")47    return {"generated_script": result}48