AlexKagan/RAG_cross_encoder

1 2 3import openai4import gradio as gr5 6from os import getenv7from typing import Any, Dict, Generator, List8 9from huggingface_hub import InferenceClient10from transformers import AutoTokenizer11 12tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")13 14temperature = 0.915top_p = 0.616repetition_penalty = 1.217 18OPENAI_KEY = getenv("OPENAI_API_KEY")19HF_TOKEN = getenv("HUGGING_FACE_HUB_TOKEN")20 21hf_client = InferenceClient(22        "mistralai/Mistral-7B-Instruct-v0.1",23        token=HF_TOKEN24        )25 26 27def format_prompt(message: str, api_kind: str):28    """29    Formats the given message using a chat template.30 31    Args:32        message (str): The user message to be formatted.33 34    Returns:35        str: Formatted message after applying the chat template.36    """37 38    # Create a list of message dictionaries with role and content39    messages: List[Dict[str, Any]] = [{'role': 'user', 'content': message}]40 41    if api_kind == "openai":42        return messages43    elif api_kind == "hf":44        return tokenizer.apply_chat_template(messages, tokenize=False)45    elif api_kind:46        raise ValueError("API is not supported")47 48 49def generate_hf(prompt: str, history: str, temperature: float = 0.9, max_new_tokens: int = 256,50             top_p: float = 0.95, repetition_penalty: float = 1.0) -> Generator[str, None, str]:51    """52    Generate a sequence of tokens based on a given prompt and history using Mistral client.53 54    Args:55        prompt (str): The initial prompt for the text generation.56        history (str): Context or history for the text generation.57        temperature (float, optional): The softmax temperature for sampling. Defaults to 0.9.58        max_new_tokens (int, optional): Maximum number of tokens to be generated. Defaults to 256.59        top_p (float, optional): Nucleus sampling probability. Defaults to 0.95.60        repetition_penalty (float, optional): Penalty for repeated tokens. Defaults to 1.0.61 62    Returns:63        Generator[str, None, str]: A generator yielding chunks of generated text.64                                   Returns a final string if an error occurs.65    """66 67    temperature = max(float(temperature), 1e-2)  # Ensure temperature isn't too low68    top_p = float(top_p)69 70    generate_kwargs = {71        'temperature': temperature,72        'max_new_tokens': max_new_tokens,73        'top_p': top_p,74        'repetition_penalty': repetition_penalty,75        'do_sample': True,76        'seed': 42,77        }78    79    formatted_prompt = format_prompt(prompt, "hf")80 81    try:82        stream = hf_client.text_generation(formatted_prompt, **generate_kwargs,83                                            stream=True, details=True, return_full_text=False)84        output = ""85        for response in stream:86            output += response.token.text87            yield output88 89    except Exception as e:90        if "Too Many Requests" in str(e):91            print("ERROR: Too many requests on Mistral client")92            gr.Warning("Unfortunately Mistral is unable to process")93            return "Unfortunately, I am not able to process your request now."94        elif "Authorization header is invalid" in str(e):95            print("Authetification error:", str(e))96            gr.Warning("Authentication error: HF token was either not provided or incorrect")97            return "Authentication error"98        else:99            print("Unhandled Exception:", str(e))100            gr.Warning("Unfortunately Mistral is unable to process")101            return "I do not know what happened, but I couldn't understand you."102 103 104def generate_openai(prompt: str, history: str, temperature: float = 0.9, max_new_tokens: int = 256,105             top_p: float = 0.95, repetition_penalty: float = 1.0) -> Generator[str, None, str]:106    """107    Generate a sequence of tokens based on a given prompt and history using Mistral client.108 109    Args:110        prompt (str): The initial prompt for the text generation.111        history (str): Context or history for the text generation.112        temperature (float, optional): The softmax temperature for sampling. Defaults to 0.9.113        max_new_tokens (int, optional): Maximum number of tokens to be generated. Defaults to 256.114        top_p (float, optional): Nucleus sampling probability. Defaults to 0.95.115        repetition_penalty (float, optional): Penalty for repeated tokens. Defaults to 1.0.116 117    Returns:118        Generator[str, None, str]: A generator yielding chunks of generated text.119                                   Returns a final string if an error occurs.120    """121 122    temperature = max(float(temperature), 1e-2)  # Ensure temperature isn't too low123    top_p = float(top_p)124    125    generate_kwargs = {126        'temperature': temperature,127        'max_tokens': max_new_tokens,128        'top_p': top_p,129        'frequency_penalty': max(-2., min(repetition_penalty, 2.)),130        }131 132    formatted_prompt = format_prompt(prompt, "openai")133 134    try:135        stream = openai.ChatCompletion.create(model="gpt-3.5-turbo-0301",136                                                messages=formatted_prompt, 137                                                **generate_kwargs, 138                                                stream=True)139        output = ""140        for chunk in stream:141            output += chunk.choices[0].delta.get("content", "")142            yield output143 144    except Exception as e:145        if "Too Many Requests" in str(e):146            print("ERROR: Too many requests on OpenAI client")147            gr.Warning("Unfortunately OpenAI is unable to process")148            return "Unfortunately, I am not able to process your request now."149        elif "You didn't provide an API key" in str(e):150            print("Authetification error:", str(e))151            gr.Warning("Authentication error: OpenAI key was either not provided or incorrect")152            return "Authentication error"153        else:154            print("Unhandled Exception:", str(e))155            gr.Warning("Unfortunately OpenAI is unable to process")156            return "I do not know what happened, but I couldn't understand you."157