| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
| import gradio as gr |
|
|
| model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" |
| model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16) |
| tokenizer = AutoTokenizer.from_pretrained(model_id) |
|
|
| generator = pipeline( |
| "text-generation", |
| model=model, |
| tokenizer=tokenizer, |
| do_sample=True, |
| top_p=0.9, |
| temperature=0.7, |
| ) |
|
|
| def respond(message, history): |
| prompt = "" |
| for user, bot in history: |
| prompt += f"<|user|>: {user}\n<|assistant|>: {bot}\n" |
| prompt += f"<|user|>: {message}\n<|assistant|>:" |
|
|
| output = generator(prompt, max_new_tokens=256)[0]["generated_text"] |
| answer = output.split("<|assistant|>:")[-1].strip() |
| return answer |
|
|
| gr.ChatInterface(respond).launch(share=True) |
|
|
