trial1 / main.py
Abhinav7's picture
Update main.py
bdf7da6 verified
raw
history blame contribute delete
No virus
1.16 kB
from fastapi import FastAPI, Form
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from ctransformers import AutoModelForCausalLM
#Model loading
llm = AutoModelForCausalLM.from_pretrained("zephyr-7b-beta.Q4_K_S.gguf",
model_type='mistral',
max_new_tokens = 256,
threads = 3,
)
#Pydantic object
class validation(BaseModel):
prompt: str
#Fast API
app = FastAPI()
# Set up CORS
origins = [
"http://localhost", # Replace with the address of your Flutter web app
"http://localhost:55345", # Add the port used by your Flutter web app
]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
#Zephyr completion
@app.post("/llm_on_cpu")
async def stream(item: validation):
system_prompt = 'Your name is Maff, a mentor who teaches people how to make money. Follow these guideline- Keep your sentences short.'
E_INST = "</s>"
user, assistant = "<|user|>", "<|assistant|>"
prompt = f"{system_prompt}{E_INST}\n{user}\n{item.prompt.strip()}{E_INST}\n{assistant}\n"
return llm(prompt)