# encoding: utf-8
import json
import openai
class Ai00:
    def __init__(self,port=65530,api_key="JUSTSECRET_KEY") :
        openai.api_base = f"http://127.0.0.1:{port}/api/oai"
        openai.api_key = api_key
        self.params = {
            "model": "model",
            "max_tokens": 4096,
            "top_p": 0,
            "temperature": 0,
            "presence_penalty": 0,
            "frequency_penalty": 1,
            "half_life": 400,
            "stop": ['\x00']
        }
        
    
    def continuation(self, message,bnf=""):
        response = openai.Completion.create(
            bnf_schema=bnf,
            model=self.params['model'],
            prompt=message,
            max_tokens=self.params['max_tokens'],
            half_life=self.params['half_life'],
            top_p=self.params['top_p'],
            temperature=self.params['temperature'],
            presence_penalty=self.params['presence_penalty'],
            frequency_penalty=self.params['frequency_penalty'],
            stop=self.params['stop']
        )
        result = response.choices[0].text
        return result
    
    def review(self,text):
        result = self.continuation(f"Question: {text}\n\nAnswer: ",
"""
<start> ::= '{' <json> '}'
<json> ::= <sex> ',' <politics> ',' <unsafe> ',' <insult>
<sex> ::= '\"sex_level\": ' <level> ',' '\"sex_misjudgment\": ' <level>
<politics> ::= '\"politics_level\": ' <level> ',' '\"politics_misjudgment\": ' <level>
<unsafe> ::= '\"unsafe_level\": ' <level> ',' '\"unsafe_misjudgment\": ' <level>
<insult> ::= '\"insult_level\": ' <level> ',' '\"insult_misjudgment\": ' <level>
<level> ::= '0' | '1' | '2' | '3'
""")
        return json.loads(result)
        

model = Ai00()
print(model.review("操你妈"))