ruslanmv/ollama-webui

1import subprocess2def check_model_exists(model_name):3    try:4        # List available models5        output = subprocess.check_output("ollama list", shell=True, stderr=subprocess.STDOUT, universal_newlines=True)6        available_models = [line.split()[0] for line in output.strip().split('\n')[1:]]7        return any(model_name in model for model in available_models)8    except subprocess.CalledProcessError as e:9        print(f"Error checking models: {e.output}")10        return False11    except Exception as e:12        print(f"An unexpected error occurred: {str(e)}")13        return False14    15    16def download_model(model_name):17    remote_models=['llama3',18    'llama3:70b',19    'phi3',20    'mistral',21    'neural-chat',22    'starling-lm',23    'codellama',24    'llama2-uncensored',25    'llava',26    'gemma:2b',27    'gemma:7b',28    'solar']29    if model_name in remote_models:30        try:31            # Download the model32            print(f"Downloading model '{model_name}'...")33            subprocess.check_call(f"ollama pull {model_name}", shell=True)34            print(f"Model '{model_name}' downloaded successfully.")35        except subprocess.CalledProcessError as e:36            print(f"Error downloading model: {e.output}")37            raise e38        except Exception as e:39            print(f"An unexpected error occurred: {str(e)}")40            raise e41    else:42        print("Not supported model currently")43 44 45def check_model(model_name):46    if not check_model_exists(model_name):47            try:48                download_model(model_name)49            except Exception as e:50                print(f"Failed to download model '{model_name}': {e}")51                return52    else:53        print("OK")54 55 56 57def make_simple_prompt(input, messages):58    """59    Create a simple prompt based on the input and messages.60    61    :param input: str, input message from the user62    :param messages: list, conversation history as a list of dictionaries containing 'role' and 'content'63    :return: str, generated prompt64    """65    if len(messages) == 1:66        prompt = f'''You are a friendly AI companion.67You should answer what the user request.68user: {input}'''69    else:70        conversation_history = '\n'.join(71            f"{message['role']}: {message['content']}" for message in reversed(messages[:-1])72        )73        prompt = f'''You are a friendly AI companion.74history: {conversation_history}.75You should answer what the user request.76user: {input}'''77 78    print(prompt)79    return prompt80 81 82def make_prompt(input, messages, model):83    """84    Create a prompt based on the input, messages, and model used.85    86    :param input: str, input message from the user87    :param messages: list, conversation history as a list of dictionaries containing 'role' and 'content'88    :param model: str, name of the model ("llama3", "mistral", or other)89    :return: str, generated prompt90    """91    if model == "llama3":92        # Special Tokens used with Meta Llama 393        BEGIN_OF_TEXT = "<|begin_of_text|>"94        EOT_ID = "<|eot_id|>"95        START_HEADER_ID = "<|start_header_id|>"96        END_HEADER_ID = "<|end_header_id|>"97    elif model == "mistral":98        # Special tokens Mistral99        BEGIN_OF_TEXT = "<s>"100        EOT_ID = "</s>"101        START_HEADER_ID = ""  # Not applicable to Mistral102        END_HEADER_ID = ""  # Not applicable to Mistral103    else:104        # No Special tokens105        BEGIN_OF_TEXT = ""106        EOT_ID = ""107        START_HEADER_ID = ""108        END_HEADER_ID = ""109 110    if len(messages) == 1:111        prompt = f'''{BEGIN_OF_TEXT}{START_HEADER_ID}system{END_HEADER_ID}112You are a friendly AI companion.113{EOT_ID}{START_HEADER_ID}user{END_HEADER_ID}114{input}115{EOT_ID}'''116    else:117        conversation_history = '\n'.join(118            f"{START_HEADER_ID}{message['role']}{END_HEADER_ID}\n{message['content']}{EOT_ID}" for message in reversed(messages[:-1])119        )120        prompt = f'''{BEGIN_OF_TEXT}{START_HEADER_ID}system{END_HEADER_ID}121You are a friendly AI companion.122history:123{conversation_history}124{EOT_ID}{START_HEADER_ID}user{END_HEADER_ID}125{input}126{EOT_ID}'''127 128    print(prompt)129    return prompt130