Spaces:
Sleeping
Sleeping
from fastapi import FastAPI | |
from pydantic import BaseModel | |
import requests | |
from ctransformers import AutoModelForCausalLM | |
llms = { | |
"tinyllama":{"name": "TheBloke/TinyLlama-1.1B-1T-OpenOrca-GGUF", "file":"tinyllama-1.1b-1t-openorca.Q4_K_M.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant <|im_end|><|im_start|>user"}, | |
"tinyllama2":{"name": "TheBloke/TinyLlama-1.1B-1T-OpenOrca-GGUF", "file":"tinyllama-1.1b-1t-openorca.Q3_K_M.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant <|im_end|><|im_start|>user"}, | |
"tinyllama3":{"name": "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF", "file":"tinyllama-1.1b-chat-v0.3.Q4_K_M.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant <|im_end|><|im_start|>user"}, | |
"tinyllama4":{"name": "TheBloke/TinyLlama-1.1B-Chat-v0.3-GGUF", "file":"tinyllama-1.1b-chat-v0.3.Q2_K.gguf", "suffix":"<|im_end|><|im_start|>assistant", "prefix":"<|im_start|>system You are a helpful assistant <|im_end|><|im_start|>user"}, | |
"solar":{"name": "TheBloke/SOLAR-10.7B-Instruct-v1.0-GGUF", "file":"solar-10.7b-instruct-v1.0.Q4_K_M.gguf", "suffix":"\n### Assistant:\n", "prefix":"### User:\n"}, | |
#"mixtral-moe":{"name": "TheBloke/Mixtral_7Bx2_MoE-GGUF", "file":"mixtral_7bx2_moe.Q4_K_M.gguf", "suffix":"\n### Assistant:\n", "prefix":"### User:\n"}, | |
#"phi2":{"name": "TheBloke/phi-2-GGUF", "file":"phi-2.Q4_K_M.gguf", "suffix":"", "prefix":""} | |
} | |
for k in llms.keys(): | |
AutoModelForCausalLM.from_pretrained(llms[k]['name'], model_file=llms[k]['file']) | |
#Pydantic object | |
class validation(BaseModel): | |
prompt: str | |
llm: str | |
#Fast API | |
app = FastAPI() | |
async def stream(item: validation): | |
prefix=llms[item.llm]['prefix'] | |
suffix=llms[item.llm]['suffix'] | |
user=""" | |
{prompt}""" | |
llm = AutoModelForCausalLM.from_pretrained(llms[item.llm]['name'], model_file=llms[item.llm]['file']) | |
prompt = f"{prefix}{user.replace('{prompt}', item.prompt)}{suffix}" | |
return llm(prompt) |