NCTCMumbai commited on
Commit
62dc511
1 Parent(s): 5186af2

Update backend/query_llm.py

Browse files
Files changed (1) hide show
  1. backend/query_llm.py +2 -1
backend/query_llm.py CHANGED
@@ -53,13 +53,14 @@ def format_prompt(message: str, api_kind: str):
53
 
54
  # Create a list of message dictionaries with role and content
55
  messages1: List[Dict[str, Any]] = [{'role': 'user', 'content': message}]
56
- messages2: List[Dict[str, Any]] = [{'role': 'user', 'content': message}]
57
 
58
  if api_kind == "openai":
59
  return messages1
60
  elif api_kind == "hf":
61
  return tokenizer.apply_chat_template(messages, tokenize=False)
62
  elif api_kind=="gemini":
 
63
  return messages2
64
  else:
65
  raise ValueError("API is not supported")
 
53
 
54
  # Create a list of message dictionaries with role and content
55
  messages1: List[Dict[str, Any]] = [{'role': 'user', 'content': message}]
56
+ messages2: List[Dict[str, Any]] = [{'role': 'user', 'parts': message}]
57
 
58
  if api_kind == "openai":
59
  return messages1
60
  elif api_kind == "hf":
61
  return tokenizer.apply_chat_template(messages, tokenize=False)
62
  elif api_kind=="gemini":
63
+ print(messages2)
64
  return messages2
65
  else:
66
  raise ValueError("API is not supported")