#pip install fastapi
#uvicorn main:app --reload
#import gradio as gr
import torch
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="TheBloke/zephyr-7B-alpha-GGUF")
#model = AutoModel.from_pretrained("TheBloke/zephyr-7B-alpha-GGUF")



pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")

# We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
messages = [
    {
        "role": "system",
        "content": "You are a Spiritual Coach who always responds in the most profound and poetic style",
    },
    {"role": "user", "content": "What is Life?"},
]
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
outputs = pipe(prompt, max_new_tokens=2560, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
print(outputs[0]["generated_text"])


@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=2560, num_return_sequences=1)
    return outputs[0]["generated_text"]