File size: 562 Bytes
61245a6 24a48da 61245a6 24a48da 61245a6 24a48da |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
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)
|