import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from huggingface_hub import hf_hub_download
import os
from huggingface_hub import login

hf_token = os.environ["HF_TOKEN"]
login(token=hf_token)


# Carica modello GRPO fine-tuned
model_name = "SelmaNajih001/GRPORagMinstral2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
pipe = pipeline(
    "text-generation",
    model=model_name,
    tokenizer=model_name,
    max_new_tokens=400,
    temperature=0.5,
    num_beams=6,
    repetition_penalty=1.5
)

# Prompt template
prompt_template = """
You are a financial market analyst.
Before making a prediction you always analyze the past, which is given by the Context below.
Answer the Question based on what happened in the past.

Please respond with:
- Chosen Stock: (name)
- Prediction: (price change)
- Explanation: (brief and clear)

Context: 
{context}

Question:
{question}
"""
prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/multi-qa-mpnet-base-dot-v1")


import tempfile
tmp_dir = tempfile.mkdtemp()

local_faiss_dir = "./faiss_index"

vectorstore = FAISS.load_local(
    local_faiss_dir, 
    embeddings
)


def analisi_finanziaria(query, k=4):
    # Recupera i documenti più rilevanti
    docs_found = vectorstore.similarity_search(query, k=k)

    # Costruisci contesto concatenando i documenti
    context = "\n".join([doc.page_content for doc in docs_found])

    # Costruisci prompt finale
    final_prompt = prompt.format(context=context, question=query)

    # Genera risposta
    result = pipe(final_prompt)[0]['generated_text']
    return result
    
iface = gr.Interface(
    fn=analisi_finanziaria,
    inputs=gr.Textbox(label="Enter event or question"),
    outputs=gr.Textbox(label="Prediction"),
    title="GRPO Financial Analyst",
    description="Enter a financial event, the GRPO model will analyze historical context and provide a prediction."
)


iface.launch()