import gradio as gr
import spaces
import torch
from peft import PeftModel
import os


from transformers import AutoModel, AutoTokenizer
from transformers.generation.utils import GenerationConfig
# 获取HF token（Spaces会自动提供）
hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")

print("=" * 50)
print("开始加载FinGPT情感分析模型...")
print("=" * 50)

model = None
tokenizer = None
device = None

try:

    # Load model directly
    model = AutoModel.from_pretrained("Go4miii/DISC-FinLLM", trust_remote_code=True, dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained("Go4miii/DISC-FinLLM", use_fast=False, trust_remote_code=True)

except Exception as e:
    print("\n" + "=" * 50)
    print("❌ 模型加载失败！")
    print(f"错误信息: {e}")
    print("=" * 50)
    raise


@spaces.GPU
def analyze_sentiment(news_text):
    """
    分析金融新闻的情感倾向
    """
    if model is None or tokenizer is None:
        return "❌ 模型未正确加载，请检查Spaces日志。"

    try:
        # 构建prompt（按照FinGPT的格式）
        prompt = f'''Instruction: What is the sentiment of this news? Please choose an answer from {{negative/neutral/positive}}
Input: {news_text}
Answer: '''

        # 编码输入
        tokens = tokenizer(
            prompt,
            return_tensors='pt',
            padding=True,
            max_length=512,
            truncation=True
        ).to(device)

        # 生成响应
        with torch.no_grad():
            res = model.generate(
                **tokens,
                max_length=512,
                pad_token_id=tokenizer.eos_token_id
            )

        # 解码输出
        res_sentence = tokenizer.decode(res[0], skip_special_tokens=True)

        # 提取答案
        if "Answer: " in res_sentence:
            sentiment = res_sentence.split("Answer: ")[1].strip()
            # 清理多余的换行和空格
            sentiment = sentiment.split('\n')[0].strip()
        else:
            sentiment = res_sentence

        return sentiment

    except Exception as e:
        return f"❌ 分析出错: {str(e)}"


# 创建Gradio界面
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 📊 FinGPT 金融新闻情感分析
        
        基于 **FinGPT/fingpt-mt_llama3-8b_lora** 模型的金融新闻情感分析工具。
        
        输入金融新闻文本，AI将分析其情感倾向：**positive（积极）** / **neutral（中性）** / **negative（消极）**
        """
    )

    with gr.Row():
        with gr.Column():
            news_input = gr.Textbox(
                label="📰 输入金融新闻",
                placeholder="粘贴或输入金融新闻内容...",
                lines=6
            )
            analyze_btn = gr.Button("🔍 分析情感", variant="primary", size="lg")

        with gr.Column():
            sentiment_output = gr.Textbox(
                label="😊 情感分析结果",
                lines=2
            )

    gr.Examples(
        examples=[
            "什么是不良资产",
        ],
        inputs=news_input,
        label="📋 示例"
    )

    # 事件处理
    analyze_btn.click(
        fn=analyze_sentiment,
        inputs=news_input,
        outputs=sentiment_output
    )

    news_input.submit(
        fn=analyze_sentiment,
        inputs=news_input,
        outputs=sentiment_output
    )

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