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

# First define a prediction function that takes in a text prompt and returns the text completion

model = pipeline("text-generation", model="zenai-org/SmolLM-prompt-generation")

def predict(prompt):
    out = model(
        prompt,
        max_length=77,    # Max length of the generated sequence
        min_length=10,    # Minimum length of the generated sequence
        do_sample=True,   # Enable sampling
        top_k=50,         # Top-k sampling
        top_p=0.95,       # Top-p sampling
        temperature=0.7,  # Control the creativity of the output
        eos_token_id=0,  # End-of-sequence token
        # pad_token_id = tokenizer.eos_token_id,
    )
    return out[0]['generated_text']

# Now create the interface
gr.Interface(fn=predict, inputs="text", outputs="text", css=".footer{display:none !important}").launch(share=True)