import spaces
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

template = """<s><|user|>
Risolvi gli indizi tra parentesi per ottenere una prima lettura, e usa la chiave di lettura per ottenere la soluzione del rebus.

{input}<|end|>
<|assistant|>"""

tokenizer = AutoTokenizer.from_pretrained("gsarti/phi3-mini-rebus-solver-fp16")
model = AutoModelForCausalLM.from_pretrained("gsarti/phi3-mini-rebus-solver-fp16")

@spaces.GPU
def solve_verbalized_rebus(example, history):
    input = template.format(input=example)
    inputs = tokenizer(input, return_tensors="pt")["input_ids"]
    outputs = model.generate(input_ids = inputs, max_new_tokens = 500, use_cache = True)
    model_generations = tokenizer.batch_decode(outputs)
    return model_generations[0]

demo = gr.ChatInterface(fn=solve_verbalized_rebus, examples=["Rebus: [Materiale espulso dai vulcani] R O [Strumento del calzolaio] [Si trovano ai lati del bacino] C I [Si ingrassano con la polenta] E I N [Contiene scorte di cibi] B [Isola in francese]\nChiave risolutiva: 1 ' 5 6 5 3 3 1 14"], title="Verbalized Rebus Solver")
demo.launch()