app.py

import gradio as gr
from transformers import pipeline
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

# Load the model
model_name = "nvidia/Llama3-ChatQA-1.5-8B"
qa_pipeline = pipeline("text-generation", model=model_name)

# FastAPI app
app = FastAPI()

class Query(BaseModel):
    inputs: str

@app.post("/predict")
async def predict(query: Query):
    response = qa_pipeline(query.inputs, max_length=250)
    return {"generated_text": response[0]["generated_text"]}

# Gradio app
def generate_answer(question):
    response = qa_pipeline(question, max_length=250)
    return response[0]["generated_text"]

iface = gr.Interface(fn=generate_answer, inputs="text", outputs="text", title="Llama3 ChatQA")

# Mount Gradio app to FastAPI
@app.get("/")
async def gradio_app():
    return gr.mount_gradio_app(app, iface)

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