from transformers import pipeline
import gradio as gr
from fastapi import FastAPI
from gradio_client import Client
from typing import Union
import uvicorn


app = FastAPI()

chatbot = pipeline(model="google/gemma-7b")

def DS_chatbot(message,history):
    conversation = chatbot(message)
    return conversation[0]['generated_text']

io = gr.ChatInterface(DS_chatbot, title=" DS Chatbot", description="Enter text to start chatting.")
@app.get("/")
async def read_main():
    return {"message": "Append /gradio to the url to see gradio the interface"
           ,"message2": "Append /hello/{any_name} to get a greeting"}

@app.get("/hello/{name}")
async def read_name(name: Union[str, None] = None):
    return { "Hey!": name}

# Mount the Gradio app onto the FastAPI app
app = gr.mount_gradio_app(app, io, path='/gradio')
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)