app.py

import gradio as gr
from PIL import Image
from ultralytics import YOLO
from transformers import pipeline
import torch

# ——— 1) Chart-pattern detector (YOLOv8) ———
pattern_model = YOLO("model.pt")

# ——— 2) Trading-Hero-LLM pipeline ———
qa_pipeline = pipeline(
    "text-generation",
    model="fuchenru/Trading-Hero-LLM",
    max_new_tokens=100,
    device=0 if torch.cuda.is_available() else -1
)

def analyze_charts(img1: Image.Image, img2: Image.Image):
    output = []

    # Chart 1
    if img1:
        res1 = pattern_model(img1)[0]
        labels1 = [res1.names[int(b.cls[0])] for b in res1.boxes]
        output.append(
            "📊 Chart 1 Patterns:\n" +
            ("\n".join(f"• {lbl}" for lbl in labels1) if labels1 else "No patterns detected.")
        )
    else:
        output.append("🖼️ Chart 1: No image uploaded.")

    # Chart 2
    if img2:
        res2 = pattern_model(img2)[0]
        labels2 = [res2.names[int(b.cls[0])] for b in res2.boxes]
        output.append(
            "\n📊 Chart 2 Patterns:\n" +
            ("\n".join(f"• {lbl}" for lbl in labels2) if labels2 else "No patterns detected.")
        )
    else:
        output.append("\n🖼️ Chart 2: No image uploaded.")

    return "\n".join(output)

def answer_question(question: str):
    if not question.strip():
        return "❌ Please enter a question."
    resp = qa_pipeline(question)[0]["generated_text"]
    return resp

# ——— Gradio UI ———
with gr.Blocks() as demo:
    gr.Markdown("## 📈 Nifty AI Trading Assistant")

    with gr.Row():
        img1 = gr.Image(type="pil", label="Upload Chart 1")
        img2 = gr.Image(type="pil", label="Upload Chart 2")

    analyze_btn = gr.Button("🔍 Analyze Charts")
    pattern_out = gr.Textbox(label="Chart Pattern Output")
    analyze_btn.click(fn=analyze_charts, inputs=[img1, img2], outputs=pattern_out)

    gr.Markdown("---")

    question = gr.Textbox(label="💬 Ask a Trading Question")
    answer_btn = gr.Button("🤖 Get LLM Response")
    llm_out = gr.Textbox(label="LLM Answer")
    answer_btn.click(fn=answer_question, inputs=question, outputs=llm_out)

demo.launch()