system_prompt =
"""
        You are a professinal few-shot text classifier. You will be given {n}-way {k}-shot examples and a set of query texts.
        Each example corresponds to one of the {n} classes. 

        Your task is to understand the class category from the supported texts, and then predict the correct class label for each query based on the provided examples of each class.

        - Each class will be labeled as "Class i: [text k]", where i is the class number and each text is an example that help you understand the class i.
        - Each query will be labeled as "Query Text: [text]" and needs to be classified into one of the {n} classes.
        - Respond with only the predicted class labels (1, 2, ..., n), in the same order as the queries.
        - Do not generate any additional text.
        - Now, take a deep breath, focus on the classification task and complete it with high accuracy.
"""
user_prompt = 
"""
        Examples:
    {Example}

        Queries:
    {Query}
    """

class FSTCClassifier(BaseModel,extra="allow"):
    n_support: int = 5
    n_query: int = 5
    n_class: int = 2
    base_url: str = "https://api.openai.com/v1"
    api_key: str = ""
    context: str = """
        You are a helpful AI assistant, skilled in classifying passages and texts into one of several categories based on examples.
    """
    """Initial prompt context."""
    model: str = "gpt-3.5-turbo"
    MAX_LENGTH: int = 16385
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        if self.api_key == "":
            self.api_key = os.environ.get("OPENAI_API_KEY")
        if os.environ.get("OPENAI_API_BASE"):
            self.base_url = os.environ.get("OPENAI_API_BASE")
        # logger.info(f"API Key: {self.api_key}")
        logger.info(f"Base URL: {self.base_url}")
        logger.info(f"Model: {self.model}")
        self.init_results()
    def init_results(self):
        self.results = {
            "prediction": [],
            "accuracy": [],
            "f1": [],
            "failed": 0,
            "success": 0
        }
        
    @pydantic.computed_field()
    @property
    def api(self) -> int:
        """Instantiates and returns an OpenAI client object."""
        client = OpenAI(api_key=self.api_key,base_url=self.base_url)
        client.timeout = 20

        return client


    def validate_response(self, resp: str):
        # split by "\n"
        predictions_1 = resp.strip().split("\n")
        # split by ","
        predictions_2 = resp.strip().split(",")
        predictions = predictions_1 if len(predictions_1) > len(predictions_2) else predictions_2 
        
        predictions = [p.strip() for p in predictions if p.strip().isdigit()]  # Ensure numeric output
        # minus 1
        predictions = list(map(lambda x: str(int(str(x).strip())-1), predictions))
        expected_count = self.n_class * self.n_query
        if len(predictions) != expected_count:
            logger.debug(f"Mismatch: Got {len(predictions)} predictions, expected {expected_count}")
            return None
        return predictions
	
	# API请求失败时自动重试
    @tenacity.retry(
        retry=tenacity.retry_if_exception_type(
            exception_types=(
                openai._exceptions.RateLimitError,
                openai._exceptions.APIConnectionError,
                ReadTimeout,
            )
        ),
        wait=tenacity.wait_random(min=2, max=10),
        stop=tenacity.stop_after_attempt(5),
        reraise=True,
    )
    def predict(self, examples:dict, queries:dict, temperature:float = 0.1,
                 top_p: float = 0.9, max_tokens: int = 100):

        prompt = FSTCPromptRequest(
            n_support=self.n_support, n_query=self.n_query, n_class=self.n_class,
            )
        prompt.format(examples, queries)
        if len(prompt) > self.MAX_LENGTH:
            self.results["failed"] += 1
            logger.error(f"Prompt length exceeds maximum limit of {self.MAX_LENGTH}")
            return None
        logger.info("Prompt:{}".format(str(prompt)))
        response = self.api.chat.completions.create(
            model=self.model,
            # prompt=str(prompt),  # Ensure
            messages=[
                {"role": "system", "content": prompt.get_system_prompt()},
                {"role": "user", "content": prompt.user_prompt},
            ],
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        logger.debug(f"Response: {response}")
        
        resp = response.choices[0].message.content.strip()
        
        pred_labels = self.validate_response(resp)
        if pred_labels is None:
            self.results["failed"] += 1
            return None
        logger.debug(f"\nPredicted Labels: {pred_labels}\nExpected Labels:   {prompt.labels}")
        
        acc = accuracy_score(prompt.labels, pred_labels)
        f1 = f1_score(prompt.labels, pred_labels, average="macro")
        
        self.results['prediction'].append(pred_labels)
        self.results['accuracy'].append(acc)
        self.results['f1'].append(f1)
        self.results["success"] += 1
        logger.debug(f"Accuracy: {acc}, F1: {f1}")
        return acc, f1
    
    def batch_predict(self, examples:List[dict], queries:List[dict], temperature:float = 0.1,
                 top_p: float = 0.9, max_tokens: int = 100, test_count:int = 1):
        self.init_results()
        for i in range(test_count):
            for example, query in tqdm.tqdm(zip(examples, queries),total=len(examples)):

                result = self.predict(example, query, temperature, top_p, max_tokens)
                if result is None:
                    logger.debug("A failure.")
                    
            logger.debug(f"Test Round {i} Results: {result}")
        result = self.caculate_score()
        return result
    
    def caculate_score(self, results=None):
        if not results:
            results = self.results
        return {
            "accuracy": sum(results['accuracy'])/len(results['accuracy']),
            "f1": sum(results['f1'])/len(results['f1']),
            "success": results["success"],
            "failed": results["failed"]
        }

def load_dataset(n_class,n_support,n_query,dataset, cv,eval=False):
    data_path = "data/{}/few_shot/{}/test.json".format(dataset, cv)
    labels_path = "data/{}/few_shot/{}/labels.test.txt".format(dataset, cv)
    fewshotdataset = FewShotDataset(
        data_path=data_path,
        labels_path=labels_path,
        n_classes=n_class,n_support=n_support,n_query=n_query,
        process_type=2, # test
        eval=eval
    )
    return fewshotdataset
def get_examples(fewshotdataset, n_example, n_query):
    examples = {}
    queries = {}
    
    # Fetch an episode from the dataset
    episode = fewshotdataset.get_episode()
    
    # Support examples (training examples)
    support = episode['xs']
    
    # Query examples (test examples)
    query = episode['xq']
    
    # Loop through support and query examples
    for label_idx, support_samples in enumerate(support):
        # For each label (first dimension), we gather the sentences from support set
        examples[label_idx] = [sample['sentence'] for sample in support_samples[:n_example]]
        
    for label_idx, query_samples in enumerate(query):
        # For each label (first dimension), we gather the sentences from query set
        queries[label_idx] = [sample['sentence'] for sample in query_samples[:n_query]]
    
    return examples, queries

n_test_episodes = 100 # 一个数据集请求100次
n_test_count = 3 # 每次请求GPT分类3次
temperature = 0.1
n_top = 0.9
max_token = 100
model_name = "gpt-3.5-turbo"

classfifier = FSTCClassifier(n_support=n_support, n_query=n_query, n_class=n_class, model=model_name)
# classfifier.predict(e,q)
results = classfifier.batch_predict(e_list,q_list,temperature,n_top,max_token,n_test_count)
results

INFO:root:Base URL: [https://aihubmix.com/v1](https://aihubmix.com/v1) 
INFO:root:Model: gpt-3.5-turbo 0it [00:00, ?it/s]
INFO:root:Predicted Labels: ['2', '1', '2', '3', '4'] Expected Labels: ['0', '1', '2', '3', '4'] INFO:root:Accuracy: 0.8, F1: 0.7333333333333333 
1it [00:01, 1.50s/it]
INFO:root:Predicted Labels: ['2', '1', '1', '1', '0'] Expected Labels: ['0', '1', '2', '3', '4'] INFO:root:Accuracy: 0.2, F1: 0.1
2it [00:02, 1.46s/it]
INFO:root:Predicted Labels: ['0', '1', '1', '3', '4'] Expected Labels: ['0', '1', '2', '3', '4'] INFO:root:Accuracy: 0.8, F1: 0.7333333333333333 
3it [00:04, 1.42s/it]
INFO:root:Predicted Labels: ['2', '1', '3', '0', '4'] Expected Labels: ['0', '1', '2', '3', '4'] INFO:root:Accuracy: 0.4, F1: 0.4 
4it [00:05, 1.46s/it]
INFO:root:Predicted Labels: ['0', '1', '2', '3', '4'] 
Expected Labels: ['0', '1', '2', '3', '4'] INFO:root:Accuracy: 1.0, F1: 1.0
INFO:root:Test Round 0 Results: {'accuracy': 0.64, 'f1': np.float64(0.5933333333333333), 'success': 5, 'failed': 0}
...
{'accuracy': 0.65, 'f1': np.float64(0.6), 'success': 20, 'failed': 0}