llm01 / app.py
ynanxiu's picture
Update app.py
7c89614 verified
raw
history blame contribute delete
349 Bytes
import gradio as gr
from transformers import pipeline
pipe = pipeline("text-generation", model="ynanxiu/olmo3-190m-zh-nano", model_kwargs={"attn_implementation": "sdpa"})
def predict(message):
output = pipe(message, max_new_tokens=256)
return output[0]["generated_text"]
gr.Interface(fn=predict, inputs="text", outputs="text").launch()