How to use open LLM available in hugging face environment and use this technique about flare HOT 2 OPEN

from langchain.schema import SystemMessage, HumanMessage, AIMessage, BaseMessage
import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from transformers import StoppingCriteria, StoppingCriteriaList, LogitsProcessor, LogitsProcessorList
from typing import List, Union, Tuple
import uuid

model_name_or_path = "TheBloke/Mixtral-8x7B-Instruct-v0.1-GPTQ"
model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,                                         
    device_map="auto",                                         
    trust_remote_code=False,                                         
    revision="main",
)

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)

class WordStoppingCriteria(StoppingCriteria):
    def __init__(self, stop_ids):
        self.stop_ids = stop_ids

    def __call__(self, input_ids, scores):
        l = len(self.stop_ids)
        return len(input_ids[0]) > l and input_ids[0][-l:].tolist() == self.stop_ids

def generate_stopping_critera(stop_word:str):
    stop_word_ids = tokenizer.encode(stop_word, add_special_tokens=False)
    return WordStoppingCriteria(stop_word_ids)


class LogitsBiasProcessor(LogitsProcessor):
    def __init__(self, token_ids_to_bias):
        self.token_ids_to_bias = token_ids_to_bias

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        for token_ids, bias in self.token_ids_to_bias.items():
            for token_id in token_ids:
                scores[:, token_id] = torch.clamp(scores[:, token_id] * bias, min=0)
        return scores

def generate_custom_logits_processor(mapping:dict):
    token_ids_to_bias = {tuple(tokenizer.encode(word, add_special_tokens=False)): bias for word, bias in mapping.items()}
    return LogitsBiasProcessor(token_ids_to_bias)

def generate(input_text, **args):
    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)

    seed = args.get('seed', None)
    if seed:
        torch.manual_seed(seed)
        
    stopping_criteria = []    
    if 'stop' in args:
        stop_words = args["stop"]
        if isinstance(stop_words, str):
            stop_words = [stop_words]
        for stop_word in stop_words:
            stopping_criteria.append(generate_stopping_critera(stop_word))
    logit_processors = []    
    if 'logit_bias' in args:
        mapping = args["logit_bias"]
        logit_processors.append(generate_custom_logits_processor(mapping))


    outputs = model.generate(
        input_ids,
        do_sample=True,
        max_length=args.get('max_tokens', 50) + len(input_ids[0]),
        repetition_penalty=args.get('repetition_penalty', 1.0),
        temperature=args.get('temperature', 1.0), 
        top_k=args.get('top_k', 0),
        top_p=args.get('top_p', 1.0),
        num_return_sequences=args.get('n', 1),
        pad_token_id=model.config.eos_token_id,
        stopping_criteria=StoppingCriteriaList(stopping_criteria),
        logits_processor=LogitsProcessorList(logit_processors),
    )

    response = {
        'id': str(uuid.uuid4()),  
        'object': 'chat.completion',
        'created': int(time.time()),  
        'model': model.__class__.__name__, 
        'choices': []
    }

    total_tokens = sum(len(output) - len(input_ids[0]) for output in outputs)

    response['usage'] = {
        'prompt_tokens': len(input_ids[0]),
        'completion_tokens': total_tokens,
        'total_tokens': total_tokens + len(input_ids[0]),
    }

    for i, output in enumerate(outputs):
        num_input_tokens = len(input_ids[0]) 

        generated_text = tokenizer.decode(output[num_input_tokens:], skip_special_tokens=True)
        
        with torch.no_grad():
            logits = model(output.unsqueeze(0))[0]
            logprobs = torch.log_softmax(logits, dim=-1)[0]

        response['choices'].append({
            'index': i,
            'message': {
                'role': 'assistant',
                'content': generated_text,
            },
            'logprobs': {
                'content': [
                    {
                        'token': tokenizer.decode([output[j]]),
                        'logprob': float(logprobs[j-1][output[j]]),
                        'bytes': [output[j].item()],
                        'top_logprobs': sorted(
                            [
                                {
                                    'token': tokenizer.decode([k]),
                                    'logprob': float(logprobs[j-1][k]),
                                    'bytes': [k.item()]
                                } for k in logprobs[j-1].topk(args.get('top_logprobs', 0)).indices
                            ],
                            key=lambda x: x['logprob'],
                            reverse=True  
                        )
                    } for j in range(num_input_tokens, len(output))
                ]
            } if args.get('logprobs', False) else None,
            'finish_reason': 'stop' if output[-1] == model.config.eos_token_id else 'length',
        })

    return response