import streamlit as st
from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

app = FastAPI()


class Input(BaseModel):
    input: str


@st.cache(allow_output_mutation=True)
def load_model():
    return pipeline("text-generation", model="klyang/MentaLLaMA-chat-7B")


model = load_model()


@app.post("/")
def generate_text(input: Input):
    result = model(input.input)
    return {"generated_text": result[0]["generated_text"]}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)