from huggingface_hub import InferenceClient
from typing import List

class InferenceChat:
    """
    A class to facilitate chat interactions using a Hugging Face model via Inference API.

    Attributes:
        model_name (str): The name of the Hugging Face model hosted on the Hub.
    """

    def __init__(self, model_name: str = "mistralai/Mixtral-8x7B-Instruct-v0.1", system_prompt: str ="You are a helpful Assistant. Let's Start Chatting: ", temperature: float = 0.2, max_new_tokens: int = 1024, top_p: float = 0.95, repetition_penalty: float = 1.0):
        """
        Initialize the InferenceChat object.

        Args:
            model_name (str): The name of the Hugging Face model hosted on the Hub.
            system_prompt (str): System Prompt to pre-define the model response.
            temperature (float): Controls the randomness of the generated text. Default is 0.2.
            max_new_tokens (int): The maximum number of tokens to generate. Default is 1024.
            top_p (float): The cumulative probability for nucleus sampling. Default is 0.95.
            repetition_penalty (float): Penalty for repeating tokens in the generated text. Default is 1.0.
        """
        self.client = InferenceClient(model_name)
        self.system_prompt = str(system_prompt)
        self.temperature = float(temperature)
        self.max_new_tokens = int(max_new_tokens)
        self.top_p = float(top_p)
        self.repetition_penalty = float(repetition_penalty)

    def format_prompt(self, message: str, history: List[tuple]) -> str:
        """
        Format the prompt by incorporating user history.

        Args:
            message (str): The user's current message.
            history (list): A list of tuples containing user prompts and bot responses.

        Returns:
            str: The formatted prompt string.
        """
        prompt = f"<s> {self.system_prompt}"
        for user_prompt, bot_response in history:
            prompt += f"[INST] {user_prompt} [/INST]"
            prompt += f" {bot_response} "
        prompt += f"[INST] {message} [/INST]"
        return prompt

    def generate(self, prompt: str, history: List[tuple]) -> str:
        """
        Generate a response based on the given prompt and chat history.

        Args:
            prompt (str): The prompt for generating the next response.
            history (list): A list of tuples containing user prompts and bot responses.

        Returns:
            str: The generated response.
        """
        if self.temperature < 1e-2:
            self.temperature = 1e-2

        generate_kwargs = dict(
            temperature=self.temperature,
            max_new_tokens=self.max_new_tokens,
            top_p=self.top_p,
            repetition_penalty=self.repetition_penalty,
            do_sample=True,
            seed=42,
        )

        formatted_prompt = self.format_prompt(prompt, history)

        stream = self.client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
        output = ""

        for response in stream:
            output += response.token.text

        return output

import gradio as gr
from nltk import word_tokenize
from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords
import nltk
import os
nltk.download('punkt')
nltk.download('wordnet')
nltk.download('stopwords')

def preprocess(sentence):
    lemmatizer = WordNetLemmatizer()
    stop_words = set(stopwords.words('english'))
    tokens = word_tokenize(sentence.lower())
    tokens = [lemmatizer.lemmatize(word) for word in tokens if word.isalnum()]
    tokens = [word for word in tokens if word not in stop_words]
    return ' '.join(tokens)

def process_input(input_text):
    history=[]
    chatbot = InferenceChat(system_prompt="Act as a Senitmental analyzer. Explain it's sentiments properly.", max_new_tokens=2048)
    sentimental_analyzer = chatbot.generate(prompt=input_text, history=history)
    print("Done")

    history=[]
    chatbot2 = InferenceChat(system_prompt="Act as a Json file, and return the Score On any 6 Parameters of what you find in the above Line.", max_new_tokens=2048)
    score_analyzer = chatbot2.generate(prompt=input_text, history=history)
    print("Done")

    history=[]
    chatbot3 = InferenceChat(system_prompt="Act as a Writing Style Teacher Analyzer. And return a simple JSON file to score it.", max_new_tokens=8096)
    grammer_analyzer = chatbot3.generate(prompt=input_text, history=history)
    print("Done")

    tokens = preprocess(input_text)
    word_count = lambda sentence: len(sentence.split())
    count_tokens = word_count(tokens)

    history=[]

    complie = f"""
    Start: >
    Initiating Text Preprocessing...
        =================================
        Progress: [-------------------------] 100%
        |                            | Percentage |
        |----------------------------|------------|
        |                            | 25/100     |
        |                            | 50/100     |
        |                            | 75/100     |
        |                            | 100/100    |
        ----------------------------

    Performing Stemming and Lemmatization...
        ===================================
        Progress: [-------------------------] 100%
        |                            | Percentage |
        |----------------------------|------------|
        |                            | 25/100     |
        |                            | 50/100     |
        |                            | 75/100     |
        |                            | 100/100    |
        ----------------------------

    Converting Text to Tokens...
        ========================
        Progress: [-------------------------] 100%
        |                            | Percentage |
        |----------------------------|------------|
        |                            | 25/100     |
        |                            | 50/100     |
        |                            | 75/100     |
        |                            | 100/100    |
        ----------------------------

    Tokens: {count_tokens}
        ========================
        {tokens}

    LLM Analyzer
        ========================
        ### Start ###

        {sentimental_analyzer}

        ### End ###

    Sentimental Score:
        ========================
        {score_analyzer}

    Grammer Analyzer
        ========================
        {grammer_analyzer}

    """

    return complie

with gr.Blocks(theme=gr.themes.Soft()) as draft_tab:
    gr.Markdown("# Mistral 8x7B Finetuned Model for Sentimental Analysis and Writing Style Analysis")

    with gr.Row():

        with gr.Column():
            essay = gr.Textbox(lines=2, label="Input Text")
            slider = gr.Slider(1, 10, value=7, label="Temperature"),
            # option = gr.Radio(["Sentimental Analysis"], label="Operations")
            # transform = gr.Radio(["Default", "Happy", "Sad", "Angry", "Unethical"], label="Tunners")
            analyse = gr.Button("Analyze")

        with gr.Column():
            output_text = gr.Textbox(lines=5, label="Output Text")
            analyse.click(fn=process_input, inputs=[essay], outputs=output_text)

draft_tab.launch(debug=True)