import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
from datasets import load_dataset
import gradio as gr

# Load the classifier pipeline for sentiment analysis (if needed)
classifier = pipeline("sentiment-analysis")

# Load model and tokenizer
model_name = "ckcl/mexc_price_model"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Use AutoModelForSequenceClassification or the appropriate model class
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Load dataset
ds = load_dataset("ckcl/BTC_USDT_dataset")

# Define the prediction function
def predict(input_text):
    # Tokenize input
    inputs = tokenizer(input_text, return_tensors="pt")
    
    # Make predictions
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Extract prediction results
    predictions = torch.argmax(outputs.logits, dim=-1)
    return str(predictions.item())

# Create Gradio interface
iface = gr.Interface(fn=predict, inputs="text", outputs="text", title="MEXC Contract Prediction", description="Predict contract prices for MEXC.")

# Launch the application
iface.launch()