Guilherme34 commited on
Commit
5359f86
1 Parent(s): f506b54

Update chat.py

Browse files
Files changed (1) hide show
  1. chat.py +1 -5
chat.py CHANGED
@@ -48,11 +48,7 @@ def generate_response(message, history):
48
  response = tokenizer.decode(output[0], skip_special_tokens=True)
49
 
50
  # Extract only the assistant's response
51
- if "<|im_start|>assistant\n" in response and "<|im_end|>" in response:
52
- assistant_response = response.split("<|im_start|>assistant\n")[1].split("<|im_end|>")[0]
53
- else:
54
- # Fallback: Return the entire generated response if special tokens are not found
55
- assistant_response = response
56
 
57
  return assistant_response
58
  except Exception as e:
 
48
  response = tokenizer.decode(output[0], skip_special_tokens=True)
49
 
50
  # Extract only the assistant's response
51
+ assistant_response = response.split("<|im_start|>assistant\n")[1].split("<|im_end|>")[0]
 
 
 
 
52
 
53
  return assistant_response
54
  except Exception as e: