File size: 999 Bytes
b5c5013 67e3e61 b5c5013 67e3e61 b5c5013 67e3e61 b5c5013 67e3e61 b5c5013 67e3e61 b5c5013 67e3e61 b5c5013 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import os
from fastapi import FastAPI
from pydantic import BaseModel
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
# Nome do modelo no Hugging Face Hub
REPO_ID = "darkshop/olama"
FILENAME = "Llama-1B-Instruct.gguf"
# Caminho local do modelo
MODEL_FILE = os.path.join(os.getcwd(), FILENAME)
# Baixa o modelo se não existir
if not os.path.exists(MODEL_FILE):
print(f"Baixando {FILENAME} do Hub...")
MODEL_FILE = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
# Inicializa o Llama
llm = Llama(model_path=MODEL_FILE)
# Cria a API FastAPI
app = FastAPI(title="Darkshop Llama API")
class Prompt(BaseModel):
text: str
max_tokens: int = 200
@app.post("/gerar")
def gerar_texto(prompt: Prompt):
result = llm(prompt.text, max_tokens=prompt.max_tokens)
return {"texto": result['text']}
@app.get("/")
def root():
return {"mensagem": "API do modelo darkshop/olama ativa!"}
# Para rodar localmente:
# uvicorn app:app --host 0.0.0.0 --port 7860
|