| 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) | |