File size: 877 Bytes
d92f1ab 442b41a d92f1ab 442b41a d92f1ab e544704 d92f1ab 442b41a e544704 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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) # ou server_name="0.0.0.0" pour local
|