Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from peft import PeftModel, PeftConfig | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer | |
| from threading import Thread | |
| config = PeftConfig.from_pretrained("cjayic/qlora-phi-1_5B-ow-fanfic") | |
| model = AutoModelForCausalLM.from_pretrained("microsoft/phi-1_5") | |
| model = PeftModel.from_pretrained(model, "cjayic/qlora-phi-1_5B-ow-fanfic") | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True) | |
| streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True) | |
| def greet(intro): | |
| inputs = tokenizer(intro, return_tensors="pt", return_attention_mask=False) | |
| streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=False, skip_special_tokens=True) | |
| generate_kwargs = dict( | |
| inputs, | |
| streamer=streamer, | |
| max_new_tokens=150, | |
| do_sample=True, | |
| #top_p=0.95, | |
| #top_k=1000, | |
| #temperature=1.0, | |
| num_beams=1, | |
| ) | |
| t = Thread(target=model.generate, kwargs=generate_kwargs) | |
| t.start() | |
| partial_message = "" | |
| for new_token in streamer: | |
| if new_token != '<': | |
| partial_message += new_token | |
| yield partial_message | |
| with gr.Blocks() as demo: | |
| inp = gr.Textbox(placeholder="Intro", value="<|startoftext|>\n# Chapter 1\n", label="Starting Text", info="Initial text that will be continued by the LLM. Use `<|startoftext|>` to generate the beginning of a chapter.") | |
| out = gr.Markdown(sanitize_html=False) | |
| btn = gr.Button() | |
| btn.click(fn=greet, inputs=[inp], outputs=[out]) | |
| demo.launch() | |