Spaces:
Running
on
L4
Running
on
L4
File size: 2,193 Bytes
a121edc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
from typing import Callable
import os
from dashscope import Generation
class QwenAgent(object):
def __init__(self,
system_prompt: str = None,
track_history: bool = True):
self.system_prompt = system_prompt
if system_prompt is None:
self.history = []
else:
self.history = [
{"role": "system", "content": system_prompt}
]
self.track_history = track_history
def basic_success_check(self, response):
if not response or not response.output or not response.output.text:
print(response)
return False
else:
return True
def run(self,
prompt: str,
top_p: float = 0.95,
temperature: float = 1.0,
seed: int = 1,
max_length: int = 1024,
max_try: int = 5,
success_check_fn: Callable = None
):
self.history.append({
"role": "user",
"content": prompt
})
success = False
try_times = 0
while try_times < max_try:
response = Generation.call(
model="qwen2-72b-instruct",
messages=self.history,
top_p=top_p,
temperature=temperature,
api_key=os.environ.get('DASHSCOPE_API_KEY'),
seed=seed,
max_length=max_length
)
if success_check_fn is None:
success_check_fn = lambda x: True
if self.basic_success_check(response) and success_check_fn(response.output.text):
response = response.output.text
self.history.append({
"role": "assistant",
"content": response
})
success = True
break
else:
try_times += 1
if not self.track_history:
if self.system_prompt is not None:
self.history = self.history[:1]
else:
self.history = []
return response, success
|