Skip to content

Commit

Permalink
Merge pull request #645 from AmitSinghShorthillsAI/main
Browse files Browse the repository at this point in the history
  • Loading branch information
zainhoda authored Sep 17, 2024
2 parents 40cab58 + a0c8007 commit 8cc87a0
Showing 1 changed file with 7 additions and 0 deletions.
7 changes: 7 additions & 0 deletions src/vanna/vllm/vllm.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,12 @@ def __init__(self, config=None):
else:
self.auth_key = None

if "temperature" in config:
self.temperature = config["temperature"]
else:
# default temperature - can be overrided using config
self.temperature = 0.7

def system_message(self, message: str) -> any:
return {"role": "system", "content": message}

Expand Down Expand Up @@ -68,6 +74,7 @@ def submit_prompt(self, prompt, **kwargs) -> str:
url = f"{self.host}/v1/chat/completions"
data = {
"model": self.model,
"temperature": self.temperature,
"stream": False,
"messages": prompt,
}
Expand Down

0 comments on commit 8cc87a0

Please sign in to comment.