import gradio as gr
from transformers import pipeline

huggingface_hub login()

# Lade Modelle von Hugging Face
llama_pipeline = pipeline("text-generation", model="meta-llama/Meta-Llama-3-8B")
#finbert_pipeline = pipeline("text-classification", model="ProsusAI/finbert")

# Funktion für die Analyse
def analyze_financial_text(model_choice, text_input):
    if not text_input:
        return "Bitte gib einen Finanztext oder eine Frage ein."
    
    if model_choice == "LLaMA 3 (Text-Generierung)":
        response = llama_pipeline(text_input, max_length=200, do_sample=True)
        return response[0]["generated_text"]
    
    #elif model_choice == "FinBERT (Sentiment-Analyse)":
        #response = finbert_pipeline(text_input)
        #return f"Ergebnis: {response[0]['label']} (Score: {response[0]['score']:.2f})"
    
    return "Ungültige Auswahl."

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Finanz-Analyse mit Sprachmodellen (LLaMA 3 & FinBERT)")

    model_choice = gr.Radio(
        #choices=["LLaMA 3 (Text-Generierung)", "FinBERT (Sentiment-Analyse)"],
        choices=["LLaMA 3 (Text-Generierung)"],
        label="Wähle ein Modell",
        value="LLaMA 3 (Text-Generierung)"
    )
    
    text_input = gr.Textbox(label="Gib hier deinen Finanztext oder eine Frage ein", lines=5)
    output = gr.Textbox(label="Modell-Response", lines=5)

    analyze_button = gr.Button("Analysieren")
    analyze_button.click(analyze_financial_text, inputs=[model_choice, text_input], outputs=output)

demo.launch()