Spaces:
Paused
Paused
File size: 1,542 Bytes
c549cf7 7745de2 c549cf7 3d4f585 7745de2 c549cf7 eaf516b c549cf7 3d4f585 c549cf7 52cffa7 c549cf7 7745de2 3d4f585 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
import transformers
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from threading import Thread
from transformers import TextIteratorStreamer
import spaces
model_name = "numfa/numfa_v2-3b"
model = AutoModelForCausalLM.from_pretrained(model_name,torch_dtype=torch.float16, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens = True)
@spaces.GPU
def generate_text(prompt, max_length, top_p, top_k):
inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
generate_kwargs = dict(
inputs,
max_length=int(max_length),top_p=float(top_p), do_sample=True, top_k=int(top_k), streamer=streamer
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
generated_text=[]
for text in streamer:
generated_text.append(text)
yield "".join(generated_text)
description = """
# Deploy your first ML app using Gradio
"""
inputs = [
gr.Textbox(label="Prompt text"),
gr.Textbox(label="max-lenth generation", value=100),
gr.Slider(0.0, 1.0, label="top-p value", value=0.95),
gr.Textbox(label="top-k", value=50,),
]
outputs = [gr.Textbox(label="Generated Text")]
demo = gr.Interface(fn=generate_text, inputs=inputs, outputs=outputs, allow_flagging=False, description=description)
demo.queue(max_size=20).launch() |