File size: 1,935 Bytes
9beebbd 3b1fc7c 9beebbd 23b172b 9beebbd 6ba4f57 9beebbd e09cd05 6ba4f57 9beebbd |
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 41 42 43 |
import os
from pathlib import Path
from llama_index.embeddings import HuggingFaceEmbedding, VoyageEmbedding
from llama_index import (load_index_from_storage, ServiceContext, StorageContext)
from llama_index import download_loader, SimpleDirectoryReader
from llama_index.retrievers import RecursiveRetriever
from llama_index.query_engine import RetrieverQueryEngine
from llama_index.llms import Anyscale
from fastapi import FastAPI
app = FastAPI()
# Define the inference model
llm = Anyscale(model="mistralai/Mistral-7B-Instruct-v0.1", api_key=os.getenv("ANYSCALE_API_KEY"))
# Define the embedding model used to embed the query.
# query_embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")
embed_model = VoyageEmbedding(model_name="voyage-01", voyage_api_key=os.getenv("VOYAGE_API_KEY"))
service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
if "index" in os.listdir():
storage_context = StorageContext.from_defaults(persist_dir=Path("./index"))
else:
dir_reader = SimpleDirectoryReader(Path('./docs'))
documents = dir_reader.load_data()
index = VectorStoreIndex.from_documents(documents, service_context=service_context)
index.storage_context.persist(Path('./index'))
storage_context = StorageContext.from_defaults(persist_dir=Path("./index"))
# Load the vector stores that were created earlier.
index = load_index_from_storage(storage_context=storage_context, service_context=service_context)
# Define query engine:
index_engine = index.as_retriever(similarity_top_k=4)
index_retriever = RecursiveRetriever("vector",retriever_dict={"vector": index_engine})
query_engine = RetrieverQueryEngine.from_args(index_retriever, service_context=service_context)
# Deploy the Ray Serve application.
@app.get("/generate")
def generate(query: str):
return str(query_engine.query(query))
if __name__ == '__main__':
uvicorn.run('main:app', reload=True) |