import gradio as gr
import os

from loaddataset import ExtractRagBenchData
from createmilvusschema import CreateMilvusDbSchema
from model import generate_response
from huggingface_hub import login
from huggingface_hub import whoami
from huggingface_hub import dataset_info


hf_token = os.getenv("HF_TOKEN")
login(hf_token)

rag_extracted_data = ExtractRagBenchData()

#invoke create milvus db function
try:
    db_collection = CreateMilvusDbSchema()
except Exception as e:
    print(f"Error creating Milvus DB schema: {e}")

print(rag_extracted_data.head(5))

def chatbot(prompt):
    return whoami()

iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Capstone Project Group 10")

if __name__ == "__main__":
    iface.launch()