# import streamlit as st
# import os
# import git
from load_model import entity_extractor 

# print(os.path.exists("./rajaatif786/TickerExtraction/entity_model2.pt"))
# import pandas as pd
# import numpy as np
# from EntityExtractor import LABEL_MAP

# #os.chdir("./TickerExtraction")

# texts=[st.text_input("Enter Text")]
# st.write(texts[0])        
# data,df=entity_extractor.input_text(texts)
# probs = entity_extractor.extract_entity_probabilities( dataset=data)
# for i in range(len(probs)):
#   prediction="Predicted Company Ticker: \n"+str(list(LABEL_MAP.keys())[list(LABEL_MAP.values()).index(np.argmax(probs[i]))])+'\n'
# st.write(prediction)
import gradio as gr
from transformers import pipeline

# Assuming your entity extraction model is loaded using a function like `load_model`
# and returns the extracted entities
def extract_entities(text):
    # Load your model here if necessary (assuming it's already loaded in the original code)
    # entity_extractor = load_model()

    # Extract entities from the text using your model
    extracted_entities = entity_extractor(text)

    return extracted_entities

# Create a Gradio interface
interface = gr.Interface(
    fn=extract_entities,
    inputs="text",
    outputs="text",
    title="Entity Extraction",
    description="Enter text to extract company tickers.",
    article="<p style='font-family:sans-serif; font-size: 16px;'>This interface uses a fine-tuned model to extract company tickers from your input text.</p>"
)

# Launch the Gradio interface
interface.launch()