| import gradio as gr |
| from transformers import ( |
| AutoModelForCausalLM, |
| AutoTokenizer, |
| pipeline, |
| logging, |
| ) |
| from peft import PeftModel |
|
|
|
|
| |
| model = AutoModelForCausalLM.from_pretrained( |
| 'TinyLlama/TinyLlama-1.1B-Chat-v1.0' |
| ) |
| model = PeftModel.from_pretrained(model, 'jezv/TinyLlama-1.1b-cmd') |
| model = model.merge_and_unload() |
|
|
| |
| tokenizer = AutoTokenizer.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0', trust_remote_code=True) |
| tokenizer.pad_token = tokenizer.eos_token |
| tokenizer.padding_side = "right" |
|
|
|
|
| def greet(prompt): |
| |
| pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_new_tokens=20) |
| messages = [ |
| { |
| "role": "system", |
| "content": "You are a chatbot assistant who helps with writing console commands.", |
| }, |
| {"role": "user", "content": prompt}, |
| ] |
| text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| result = pipe(text) |
| return result[0]['generated_text'] |
|
|
| demo = gr.Interface( |
| fn=greet, |
| inputs=["text"], |
| outputs=["text"], |
| ) |
|
|
| demo.launch() |
