from typing import *
import torch
from transformers import pipeline
import spaces
from .common import Grader, device

model_name = "JacobLinCool/mistral-7b-ielts-evaluator-safetensors"
pipe = pipeline("text-generation", model_name, device=device)

system = """In this task, you are required to evaluate an IELTS Writing Task 2 essay. Consider the following four criteria and provide a detailed assessment for each, along with a suggested band score:

## Task Achievement:
- Evaluate how well the candidate has addressed the given task.
- Assess the clarity and coherence of the response in presenting ideas.
- Identify if the candidate has fully covered all parts of the task and supported arguments appropriately.
- Suggested Band Score (Task Achievement): [Insert Score]

## Coherence and Cohesion:
- Assess the overall organization and structure of the essay.
- Evaluate the use of linking devices to connect ideas and paragraphs.
- Identify if there is a logical flow of information.
- Suggested Band Score (Coherence and Cohesion): [Insert Score]

## Lexical Resource (Vocabulary):
- Examine the range and accuracy of vocabulary used in the essay.
- Point out specific mistakes in vocabulary, such as inaccuracies or overuse of certain words and Suggest modified versions or alternatives for the identified mistakes. [list of mistakes and rectify]
- Assess the appropriateness of vocabulary for the given context.
- Suggested Band Score (Lexical Resource): [Insert Score]

## Grammatical Range and Accuracy:
- Evaluate the variety and complexity of sentence structures.
- Point out specific grammatical errors, such as incorrect verb forms or sentence construction and Suggest modified versions or corrections for the identified mistakes. [list of mistakes and rectify]
- Examine the use of punctuation and sentence formation.
- Suggested Band Score (Grammatical Range and Accuracy): [Insert Score]

## Overall Band Score:

- Provide an overall band score for the essay, considering the holistic performance across all criteria.
- Consider the synergy of the essay in meeting the task requirements cohesively.
- Suggested Overall Band Score: [Insert Score]

## Feedback and Additional Comments:
- Provide constructive feedback highlighting specific strengths and areas for improvement.
- Suggest strategies for enhancement in weaker areas."""


class Mistral_7b_IELTS_Evaluator(Grader):
    def info(self) -> str:
        return "Safetensors version of [chillies/mistral-7b-ielts-evaluator-q4](https://huggingface.co/chillies/mistral-7b-ielts-evaluator-q4)"

    @spaces.GPU(duration=120)
    @torch.no_grad()
    def grade(self, question: str, answer: str) -> Tuple[float, str]:
        text = f"""{system}

## Prompt:
{question}

## Essay:
{answer}

## Evaluation:
"""

        outputs = pipe(text, max_length=2048)
        comment = outputs[0]["generated_text"].split("## Evaluation:")[1].strip()
        overall_score = float(
            comment.split("Suggested Overall Band Score: ")[1].split("\n")[0] or 0.0
        )

        return overall_score, comment