import torch
from transformers import AutoAdapterModel, AutoTokenizer
from datasets import load_dataset
import gradio as gr

# 加载模型和分词器
model_name = "ckcl/mexc_price_model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoAdapterModel.from_pretrained(model_name)
model.load_adapter(model_name, set_active=True)

# 加载数据集
ds = load_dataset("ckcl/BTC_USDT_dataset")

# 定义预测函数
def predict(input_text):
    # 处理输入
    inputs = tokenizer(input_text, return_tensors="pt")
    
    # 进行预测
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 获取预测结果
    predictions = torch.argmax(outputs.logits, dim=-1)
    return str(predictions.item())

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

# 启动应用
iface.launch()