import gradio as gr
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import torch
import theme
import chatbot

theme = theme.Theme()

# Cell 1: Image Classification Model
image_pipeline = pipeline(task="image-classification", model="guillen/vit-basura-test1")

def predict_image(input_img):
    predictions = image_pipeline(input_img)
    return {p["label"]: p["score"] for p in predictions} 

image_gradio_app = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(label="Image", sources=['upload', 'webcam'], type="pil"),
    outputs=[gr.Label(label="Result")],
    title="Green Greta",
    theme=theme
)

# Cell 2: Chatbot Model

def qa_response(user_message, chat_history, context):
    response = qa_chain.predict(user_message, chat_history, context=context)
    return response

chatbot_gradio_app = gr.ChatInterface(
    fn=qa_response,
    title="Green Greta",
    theme=theme
)

# Combine both interfaces into a single app
gr.TabbedInterface(
    [image_gradio_app, chatbot_gradio_app],
    tab_names=["Green Greta Image Classification","Green Greta Chat"],
    theme=theme
).launch()