TrgTuan10's picture
add app.py
406df24
raw
history blame
937 Bytes
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)