import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "TheFinAI/Fino1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)


@spaces.GPU
def chat(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
    output = model.generate(**inputs, max_new_tokens=200)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response


demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(label="输入问题", placeholder="例如: What is the macd ?"),
    outputs=gr.Textbox(label="回答"),
    title="FinChatBox",
    description="使用 TheFinAI 模型进行对话"
)
demo.launch()