| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| import torch | |
| model = f"tiiuae/falcon-7b" | |
| tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model, | |
| torch_dtype=torch.bfloat16, | |
| device_map="auto", | |
| load_in_8bit=True, | |
| trust_remote_code=True | |
| ) | |
| def greet(prompt): | |
| inputs = tokenizer(prompt, return_tensors="pt").to("cuda") | |
| v = model.generate( | |
| input_ids=inputs["input_ids"], | |
| attention_mask=inputs["attention_mask"], | |
| do_sample=True, | |
| temperature=0.6, | |
| top_p=0.9, | |
| max_new_tokens=50, | |
| ) | |
| return tokenizer.decode(v[0].to("cpu")) | |
| iface = gr.Interface(fn=greet, inputs="text", outputs="text") | |
| iface.launch() | |