import gradio as gr
from transformers import pipeline
import requests
import os

start_token = "<s>"
start_instruction_token = "[INST] "
end_instruction_token = " [/INST]"
system_prompt = "Sei un assistente utile ed affidabile. Rispondi in maniera adeguata alla domanda seguente:\n"
start_completion = "\nRisposta:"

API_URL = "https://cyk11dj2ce5ybyjq.us-east-1.aws.endpoints.huggingface.cloud"
token = "Bearer " + os.getenv("ITACA_TOKEN")

headers = {
    "Accept": "application/json",
    "Authorization": token,
    "Content-Type": "application/json"
}


def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()


def predict(message, history):
    new_message = start_token + start_instruction_token + system_prompt + message + end_instruction_token + start_completion
    print(new_message)
    output = query({
        "inputs": message,
        "parameters": {
    		"max_new_tokens": 512,
    		"return_full_text": False
    	}
    })
    return output[0]["generated_text"]
    #return """Gli ISA, ovvero gli Indici Sintetici di Affidabilità, sono strumenti utilizzati dall'Agenzia delle Entrate per valutare l'affidabilità fiscale dei contribuenti. Questi indici considerano diversi aspetti come il volume delle vendite, i costi del personale, le spese per l'acquisto di beni e servizi, e altri fattori che possono influire sulla regolarità fiscale. Gli ISA aiutano a determinare l'affidabilità fiscale complessiva di un'impresa o di un professionista."""

iface = gr.ChatInterface(predict)
iface.launch()