llamaindex_demo / app.py
BinZhang
dftmsg
3e029de
raw
history blame
1.7 kB
import streamlit as st
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.settings import Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai_like import OpenAILike
import os
from dotenv import load_dotenv
# 加载环境变量
load_dotenv()
# 设置 API 参数
base_url = "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/"
api_key = os.getenv("MY_API_KEY")
model = "internlm2.5-latest"
# 初始化 LLM
llm = OpenAILike(model=model, api_base=base_url, api_key=api_key, is_chat_model=True)
# 初始化一个 HuggingFaceEmbedding 对象,用于将文本转换为向量表示
embed_model = HuggingFaceEmbedding(
model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)
# 将创建的嵌入模型赋值给全局设置的 embed_model 属性
Settings.embed_model = embed_model
# 将创建的 LLM 赋值给全局设置的 llm 属性
Settings.llm = llm
# 从指定目录读取所有文档,并加载数据到内存中
documents = SimpleDirectoryReader("./data").load_data()
# 创建一个 VectorStoreIndex,并使用之前加载的文档来构建索引
index = VectorStoreIndex.from_documents(documents)
# 创建一个查询引擎,这个引擎可以接收查询并返回相关文档的响应
query_engine = index.as_query_engine()
# 设置页面标题
st.title("LlamaIndex Chat")
# 创建一个文本输入框供用户输入问题
user_input = st.text_input("请输入你的问题:")
# 按钮用于提交问题
if st.button("发送"):
# 使用查询引擎获取回复
response = query_engine.query(user_input)
# 显示回复
st.write(f"回复: {response}")