main.py

#pip install fastapi
#uvicorn main:app --reload
#import gradio as gr

from transformers import pipeline
from fastapi import FastAPI

app = FastAPI()

#generator = pipeline('text-generation',model='gpt2')
#generator = pipeline('text-generation',model='Open-Orca/Mistral-7B-OpenOrca')
generator = pipeline("text-generation", model="lmsys/vicuna-7b-v1.5")


@app.get("/")
async def root():
    return {"message": "Hello World"}
    #return generator('What is love',max_length=100, num_return_sequences=1)

@app.post("/predict")
async def root(text):
    #return {"message": "Hello World"}
    return generator(text,max_length=100, num_return_sequences=1)