import gradio as gr
from textPreprocessing import text2prompt
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

model_id = 'ai-forever/ruT5-large'

"""bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="fp4",
    bnb_4bit_compute_dtype=torch.bfloat16
)"""

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    # quantization_config=bnb_config # Необходимо раскомментить при досутпе к GPU
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

def predict(input_text, t, m):
    """
        Вывести финансовую рекомендацию на основе:
        input_text: str 
            - Контекст в виде новости из области экономики
        t: tokenizer
            - Токенизатор для модели
        m: model
            - Instruct-based модель
    """
    prompt = text2prompt(input_text)
    inputs = tokenizer(prompt, return_tensors="pt")
    generate_ids = model.generate(inputs.input_ids, max_new_tokens=128)
    answer = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    return answer.replace(prompt, "")

gradio_app = gr.Interface(
    predict,
    inputs=gr.Textbox(
        label="Входная новость", container=True,
        lines=8, placeholder="Акции кредитного банка \"X\" обрушились в цене из-за дефолта по ипотечным кредитам"
    ),
    outputs=[gr.Label(label="Финансовая рекомендация на основе новости:")],
    title="Finam Finetuned Mistral Instruct (FFMI)",
)

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