import gradio as gr
from transformers import pipeline

def stockmarket_qa(question):
    model_name = "AventIQ-AI/t5-stockmarket-qa-chatbot"
    print("Loading model...")
    generator = pipeline("text2text-generation", model=model_name, device=0)  # Use GPU if available
    
    # Correcting the input format based on the model's README
    instruction = "question: "  # If required by the model
    input_text = f"{instruction}{question}"
    
    print(f"Processing input: {input_text}")
    response = generator(input_text, max_new_tokens=200)[0]['generated_text']
    
    print(f"Response: {response}")
    return response

iface = gr.Interface(
    fn=stockmarket_qa,
    inputs=gr.Textbox(label="📈 Ask a Stock Market Question", placeholder="Enter your query about the stock market...", lines=3),
    outputs=gr.Textbox(label="💡 Expert Answer", interactive=True),
    title="📊 AI-Powered Stock Market Q&A Chatbot",
    description="🤖 Enter a stock market-related question, and the AI chatbot will generate a detailed response based on financial knowledge.",
    theme="compact",
    allow_flagging="never",
    examples=[
        ["What are the best stocks to invest in right now?"],
        ["How does inflation affect stock prices?"],
        ["What is technical analysis in stock trading?"],
        ["Explain the difference between a bull and bear market."],
        ["How do interest rates impact the stock market?"]
    ],
)

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