harsh-manvar commited on
Commit
721d85a
1 Parent(s): c08ef2a

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +7 -6
model.py CHANGED
@@ -9,14 +9,14 @@ from transformers import TextIteratorStreamer, AutoTokenizer
9
  logging.set_verbosity_info()
10
  logger = logging.get_logger("transformers")
11
 
12
- config = {'max_new_tokens': 256, 'repetition_penalty': 1.1,
13
- 'temperature': 0.1, 'stream': True}
14
- model_id = 'TheBloke/Llama-2-7B-Chat-GGML'
15
  device = "cpu"
16
 
17
 
18
- model = AutoModelForCausalLM.from_pretrained(model_id, model_type="llama", lib='avx2', hf=True)
19
- tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-2-7b-chat-hf')
20
 
21
  def get_prompt(message: str, chat_history: list[tuple[str, str]],
22
  system_prompt: str) -> str:
@@ -75,5 +75,6 @@ def run(message: str,
75
 
76
  outputs = []
77
  for text in streamer:
 
78
  outputs.append(text)
79
- yield ''.join(outputs)
 
9
  logging.set_verbosity_info()
10
  logger = logging.get_logger("transformers")
11
 
12
+ config = {"max_new_tokens": 256, "repetition_penalty": 1.1,
13
+ "temperature": 0.1, "stream": True}
14
+ model_id = "TheBloke/Llama-2-7B-Chat-GGML"
15
  device = "cpu"
16
 
17
 
18
+ model = AutoModelForCausalLM.from_pretrained(model_id, model_type="llama", lib="avx2", hf=True)
19
+ tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
20
 
21
  def get_prompt(message: str, chat_history: list[tuple[str, str]],
22
  system_prompt: str) -> str:
 
75
 
76
  outputs = []
77
  for text in streamer:
78
+ logger.info("outputs", outputs)
79
  outputs.append(text)
80
+ yield "".join(outputs)