| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
| import gradio as gr |
|
|
| model_name = "mistralai/Mistral-7B-v0.1" |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_name, |
| device_map="auto", |
| load_in_4bit=True |
| ) |
|
|
| pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) |
|
|
| def generate(text): |
| outputs = pipe(text, max_length=100, do_sample=True) |
| return outputs[0]["generated_text"] |
|
|
| demo = gr.Interface(fn=generate, inputs="text", outputs="text") |
| demo.launch() |
