Views
No views yet
1from typing import Dict, List, Optional
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, TextStreamer
4
5model_name_or_path = "bofenghuang/vigogne-stablelm-3b-4e1t-chat"
6tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="right", use_fast=False)
7model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
8
9streamer = TextStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
10
11
12def chat(
13 query: str,
14 history: Optional[List[Dict]] = None,
15 temperature: float = 0.7,
16 top_p: float = 1.0,
17 top_k: float = 0,
18 repetition_penalty: float = 1.1,
19 max_new_tokens: int = 1024,
20 **kwargs,
21):
22 if history is None:
23 history = []
24
25 history.append({"role": "user", "content": query})
26
27 input_ids = tokenizer.apply_chat_template(history, return_tensors="pt").to(model.device)
28 input_length = input_ids.shape[1]
29
30 generated_outputs = model.generate(
31 input_ids=input_ids,
32 generation_config=GenerationConfig(
33 temperature=temperature,
34 do_sample=temperature > 0.0,
35 top_p=top_p,
36 top_k=top_k,
37 repetition_penalty=repetition_penalty,
38 max_new_tokens=max_new_tokens,
39 pad_token_id=tokenizer.eos_token_id,
40 **kwargs,
41 ),
42 streamer=streamer,
43 return_dict_in_generate=True,
44 )
45
46 generated_tokens = generated_outputs.sequences[0, input_length:]
47 generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
48
49 history.append({"role": "assistant", "content": generated_text})
50
51 return generated_text, history
52
53
54# 1st round
55response, history = chat("Un escargot parcourt 100 mètres en 5 heures. Quelle est sa vitesse ?", history=None)