import gradio as gr
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

model_id = "amielitos/text-to-markdown-t5"
tokenizer = T5Tokenizer.from_pretrained(model_id)
model = T5ForConditionalGeneration.from_pretrained(model_id)

def predict(input_text):
    # Match the prefix used in training
    prompt = f"format md: {input_text}"
    inputs = tokenizer(prompt, return_tensors="pt").input_ids
    
    # beam_search + penalty = better guessing
    outputs = model.generate(
        inputs, 
        max_length=256, 
        num_beams=5, 
        repetition_penalty=3.0, 
        early_stopping=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(fn=predict, inputs="text", outputs="markdown")
demo.launch()