content: # 安装 pdf 解析库
  # !pip install pdfminer.six
  
  from pdfminer.high_level import extract_pages
  from pdfminer.layout import LTTextContainer
  
  def extract_text_from_pdf(filename, page_numbers=None, min_line_length=1):
      '''从 PDF 文件中（按指定页码）提取文字'''
      paragraphs = []
      buffer = ''
      full_text = ''
      # 提取全部文本
      for i, page_layout in enumerate(extract_pages(filename)):
          # 如果指定了页码范围，跳过范围外的页
          if page_numbers is not None and i not in page_numbers:
              continue
          for element in page_layout:
              if isinstance(element, LTTextContainer):
                  full_text += element.get_text() + '\n'
      # 按空行分隔，将文本重新组织成段落
      lines = full_text.split('\n')
      for text in lines:
          if len(text) >= min_line_length:
              buffer += (' '+text) if not text.endswith('-') else text.strip('-')
          elif buffer:
              paragraphs.append(buffer)
              buffer = ''
      if buffer:
          paragraphs.append(buffer)
      return paragraphs
    
  paragraphs = extract_text_from_pdf("llama2.pdf", min_line_length=10)
  
  for para in paragraphs[:4]:
      print(para+"\n")
      
      
  from elasticsearch7 import Elasticsearch, helpers
  from nltk.stem import PorterStemmer
  from nltk.tokenize import word_tokenize
  from nltk.corpus import stopwords
  import nltk
  import re
  
  import warnings
  warnings.simplefilter("ignore")  # 屏蔽 ES 的一些Warnings
  
  # 实验室平台已经内置
  # nltk.download('punkt')  # 英文切词、词根、切句等方法
  # nltk.download('stopwords')  # 英文停用词库
  
  
  
  def to_keywords(input_string):
      '''（英文）文本只保留关键字'''
      # 使用正则表达式替换所有非字母数字的字符为空格
      no_symbols = re.sub(r'[^a-zA-Z0-9\s]', ' ', input_string)
      word_tokens = word_tokenize(no_symbols)
      # 加载停用词表
      stop_words = set(stopwords.words('english'))
      ps = PorterStemmer()
      # 去停用词，取词根
      filtered_sentence = [ps.stem(w)
                           for w in word_tokens if not w.lower() in stop_words]
      return ' '.join(filtered_sentence)
  
  
  import os, time
  
  # 引入配置文件
  ELASTICSEARCH_BASE_URL = os.getenv('ELASTICSEARCH_BASE_URL')
  ELASTICSEARCH_PASSWORD = os.getenv('ELASTICSEARCH_PASSWORD')
  ELASTICSEARCH_NAME= os.getenv('ELASTICSEARCH_NAME')
  
  # tips: 如果想在本地运行，请在下面一行 print(ELASTICSEARCH_BASE_URL) 获取真实的配置
  
  # 1. 创建Elasticsearch连接
  es = Elasticsearch(
      hosts=[ELASTICSEARCH_BASE_URL],  # 服务地址与端口
      http_auth=(ELASTICSEARCH_NAME, ELASTICSEARCH_PASSWORD),  # 用户名，密码
  )
  
  # 2. 定义索引名称
  index_name = "teacher_demo_index"
  
  # 3. 如果索引已存在，删除它（仅供演示，实际应用时不需要这步）
  if es.indices.exists(index=index_name):
      es.indices.delete(index=index_name)
  
  # 4. 创建索引
  es.indices.create(index=index_name)
  
  # 5. 灌库指令
  actions = [
      {
          "_index": index_name,
          "_source": {
              "keywords": to_keywords(para),
              "text": para
          }
      }
      for para in paragraphs
  ]
  
  # 6. 文本灌库
  helpers.bulk(es, actions)
  
  # 灌库是异步的
  time.sleep(2)
  
  
  def search(query_string, top_n=3):
      # ES 的查询语言
      search_query = {
          "match": {
              "keywords": to_keywords(query_string)
          }
      }
      res = es.search(index=index_name, query=search_query, size=top_n)
      return [hit["_source"]["text"] for hit in res["hits"]["hits"]]
  
  results = search("how many parameters does llama 2 have?", 2)
  for r in results:
      print(r+"\n")
      
  
  
  from openai import OpenAI
  import os
  # 加载环境变量
  from dotenv import load_dotenv, find_dotenv
  _ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY
  
  client = OpenAI()
  
  
  def get_completion(prompt, model="gpt-3.5-turbo"):
      '''封装 openai 接口'''
      messages = [{"role": "user", "content": prompt}]
      response = client.chat.completions.create(
          model=model,
          messages=messages,
          temperature=0,  # 模型输出的随机性，0 表示随机性最小
      )
      return response.choices[0].message.content
    
  def build_prompt(prompt_template, **kwargs):
      '''将 Prompt 模板赋值'''
      inputs = {}
      for k, v in kwargs.items():
          if isinstance(v, list) and all(isinstance(elem, str) for elem in v):
              val = '\n\n'.join(v)
          else:
              val = v
          inputs[k] = val
      return prompt_template.format(**inputs) # format的关键词参数的用法
    
  prompt_template = """
  你是一个问答机器人。
  你的任务是根据下述给定的已知信息回答用户问题。
  
  已知信息:
  {context}
  
  用户问：
  {query}
  
  如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复"我无法回答您的问题"。
  请不要输出已知信息中不包含的信息或答案。
  请用中文回答用户问题。
  """
  
  
  user_query = "how many parameters does llama 2 have?"
  
  # 1. 检索
  search_results = search(user_query, 2)
  
  # 2. 构建 Prompt
  prompt = build_prompt(prompt_template, context=search_results, query=user_query)
  print("===Prompt===")
  print(prompt)
  
  # 3. 调用 LLM
  response = get_completion(prompt)
  
  print("===回复===")
  print(response)
  
  
  import numpy as np
  from numpy import dot
  from numpy.linalg import norm
  
  def cos_sim(a, b):
      '''余弦距离 -- 越大越相似'''
      return dot(a, b)/(norm(a)*norm(b))
  
  
  def l2(a, b):
      '''欧氏距离 -- 越小越相似'''
      x = np.asarray(a)-np.asarray(b)
      return norm(x)
    
  def get_embeddings(texts, model="text-embedding-ada-002", dimensions=None):
      '''封装 OpenAI 的 Embedding 模型接口'''
      if model == "text-embedding-ada-002":
          dimensions = None
      if dimensions:
          data = client.embeddings.create(
              input=texts, model=model, dimensions=dimensions).data
      else:
          data = client.embeddings.create(input=texts, model=model).data
      return [x.embedding for x in data]
    
  test_query = ["测试文本"]
  vec = get_embeddings(test_query)[0]
  print(f"Total dimension: {len(vec)}")
  print(f"First 10 elements: {vec[:10]}")
  
  
  # query = "国际争端"
  
  # 且能支持跨语言
  query = "global conflicts"
  
  documents = [
      "联合国就苏丹达尔富尔地区大规模暴力事件发出警告",
      "土耳其、芬兰、瑞典与北约代表将继续就瑞典“入约”问题进行谈判",
      "日本岐阜市陆上自卫队射击场内发生枪击事件 3人受伤",
      "国家游泳中心（水立方）：恢复游泳、嬉水乐园等水上项目运营",
      "我国首次在空间站开展舱外辐射生物学暴露实验",
  ]
  
  query_vec = get_embeddings([query])[0]
  doc_vecs = get_embeddings(documents)
  
  print("Query与自己的余弦距离: {:.2f}".format(cos_sim(query_vec, query_vec)))
  print("Query与Documents的余弦距离:")
  for vec in doc_vecs:
      print(cos_sim(query_vec, vec))
  
  print()
  
  print("Query与自己的欧氏距离: {:.2f}".format(l2(query_vec, query_vec)))
  print("Query与Documents的欧氏距离:")
  for vec in doc_vecs:
      print(l2(query_vec, vec))
      
      
  # !pip install chromadb
  # 为了演示方便，我们只取两页（第一章）
  paragraphs = extract_text_from_pdf(
      "llama2.pdf",
      page_numbers=[2, 3],
      min_line_length=10
  )
  
  import chromadb
  from chromadb.config import Settings
  
  
  class MyVectorDBConnector:
      def __init__(self, collection_name, embedding_fn):
          chroma_client = chromadb.Client(Settings(allow_reset=True))
  
          # 为了演示，实际不需要每次 reset()
          chroma_client.reset()
  
          # 创建一个 collection
          self.collection = chroma_client.get_or_create_collection(
              name=collection_name)
          self.embedding_fn = embedding_fn
  
      def add_documents(self, documents):
          '''向 collection 中添加文档与向量'''
          self.collection.add(
              embeddings=self.embedding_fn(documents),  # 每个文档的向量
              documents=documents,  # 文档的原文
              ids=[f"id{i}" for i in range(len(documents))]  # 每个文档的 id
          )
  
      def search(self, query, top_n):
          '''检索向量数据库'''
          results = self.collection.query(
              query_embeddings=self.embedding_fn([query]),
              n_results=top_n
          )
          return results
        
  # 创建一个向量数据库对象
  vector_db = MyVectorDBConnector("demo", get_embeddings)
  # 向向量数据库中添加文档
  vector_db.add_documents(paragraphs)
  
  user_query = "Llama 2有多少参数"
  results = vector_db.search(user_query, 2)
  
  for para in results['documents'][0]:
      print(para+"\n")
      
      
      
  class RAG_Bot:
      def __init__(self, vector_db, llm_api, n_results=2):
          self.vector_db = vector_db
          self.llm_api = llm_api
          self.n_results = n_results
  
      def chat(self, user_query):
          # 1. 检索
          search_results = self.vector_db.search(user_query, self.n_results)
  
          # 2. 构建 Prompt
          prompt = build_prompt(
              prompt_template, context=search_results['documents'][0], query=user_query)
  
          # 3. 调用 LLM
          response = self.llm_api(prompt)
          return response
        
  # 创建一个RAG机器人
  bot = RAG_Bot(
      vector_db,
      llm_api=get_completion
  )
  
  user_query = "llama 2有多少参数?"
  
  response = bot.chat(user_query)
  
  print(response)
  
  
  
  
  # 调用文心千帆 调用 BGE Embedding 接口
  
  import json
  import requests
  import os
  
  # 通过鉴权接口获取 access token
  
  def get_access_token():
      """
      使用 AK，SK 生成鉴权签名（Access Token）
      :return: access_token，或是None(如果错误)
      """
      url = "https://aip.baidubce.com/oauth/2.0/token"
      params = {
          "grant_type": "client_credentials",
          "client_id": os.getenv('ERNIE_CLIENT_ID'),
          "client_secret": os.getenv('ERNIE_CLIENT_SECRET')
      }
  
      return str(requests.post(url, params=params).json().get("access_token"))
  
  def get_embeddings_bge(prompts):
      url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/bge_large_en?access_token=" + get_access_token()
      payload = json.dumps({
          "input": prompts
      })
      headers = {'Content-Type': 'application/json'}
  
      response = requests.request(
          "POST", url, headers=headers, data=payload).json()
      data = response["data"]
      return [x["embedding"] for x in data]
  
  
  # 调用文心4.0对话接口
  def get_completion_ernie(prompt):
  
      url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + get_access_token()
      payload = json.dumps({
          "messages": [
              {
                  "role": "user",
                  "content": prompt
              }
          ]
      })
  
      headers = {'Content-Type': 'application/json'}
  
      response = requests.request(
          "POST", url, headers=headers, data=payload).json()
  
      return response["result"]
    
    
  model = "text-embedding-3-large"
  dimensions = 128
  
  query = "国际争端"
  
  # 且能支持跨语言
  # query = "global conflicts"
  
  documents = [
      "联合国就苏丹达尔富尔地区大规模暴力事件发出警告",
      "土耳其、芬兰、瑞典与北约代表将继续就瑞典“入约”问题进行谈判",
      "日本岐阜市陆上自卫队射击场内发生枪击事件 3人受伤",
      "国家游泳中心（水立方）：恢复游泳、嬉水乐园等水上项目运营",
      "我国首次在空间站开展舱外辐射生物学暴露实验",
  ]
  
  query_vec = get_embeddings([query], model=model, dimensions=dimensions)[0]
  doc_vecs = get_embeddings(documents, model=model, dimensions=dimensions)
  
  print("向量维度: {}".format(len(query_vec)))
  
  print()
  
  print("Query与Documents的余弦距离:")
  for vec in doc_vecs:
      print(cos_sim(query_vec, vec))
  
  print()
  
  print("Query与Documents的欧氏距离:")
  for vec in doc_vecs:
      print(l2(query_vec, vec))
  
  
  
  
  
  
  
  
  
  
      
  # 创建一个向量数据库对象
  vector_db = MyVectorDBConnector("demo_text_split", get_embeddings)
  # 向向量数据库中添加文档
  vector_db.add_documents(paragraphs)
  
  # 创建一个RAG机器人
  bot = RAG_Bot(
      vector_db,
      llm_api=get_completion
  )
  
  # user_query = "llama 2有商用许可协议吗"
  user_query="llama 2 chat有多少参数"
  search_results = vector_db.search(user_query, 2)
  
  for doc in search_results['documents'][0]:
      print(doc+"\n")
  
  print("====回复====")
  bot.chat(user_query)
  
  
  from nltk.tokenize import sent_tokenize
  import json
  
  
  def split_text(paragraphs, chunk_size=300, overlap_size=100):
      '''按指定 chunk_size 和 overlap_size 交叠割文本'''
      sentences = [s.strip() for p in paragraphs for s in sent_tokenize(p)]
      chunks = []
      i = 0
      while i < len(sentences):
          chunk = sentences[i]
          overlap = ''
          prev_len = 0
          prev = i - 1
          # 向前计算重叠部分
          while prev >= 0 and len(sentences[prev])+len(overlap) <= overlap_size:
              overlap = sentences[prev] + ' ' + overlap
              prev -= 1
          chunk = overlap+chunk
          next = i + 1
          # 向后计算当前chunk
          while next < len(sentences) and len(sentences[next])+len(chunk) <= chunk_size:
              chunk = chunk + ' ' + sentences[next]
              next += 1
          chunks.append(chunk)
          i = next
      return chunks
    
    
  # !pip install sentence_transformers
  
  from sentence_transformers import CrossEncoder
  
  # model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', max_length=512) # 英文，模型较小
  model = CrossEncoder('BAAI/bge-reranker-large', max_length=512) # 多语言，国产，模型较大
  
  user_query = "how safe is llama 2"
  # user_query = "llama 2安全性如何"
  scores = model.predict([(user_query, doc)
                         for doc in search_results['documents'][0]])
  # 按得分排序
  sorted_list = sorted(
      zip(scores, search_results['documents'][0]), key=lambda x: x[0], reverse=True)
  for score, doc in sorted_list:
      print(f"{score}\t{doc}\n")
      
      
  def rrf(ranks, k=1):
      ret = {}
      # 遍历每次的排序结果
      for rank in ranks:
          # 遍历排序中每个元素
          for id, val in rank.items():
              if id not in ret:
                  ret[id] = {"score": 0, "text": val["text"]}
              # 计算 RRF 得分
              ret[id]["score"] += 1.0/(k+val["rank"])
      # 按 RRF 得分排序，并返回
      return dict(sorted(ret.items(), key=lambda item: item[1]["score"], reverse=True))
    
  import json
  
  # 融合两次检索的排序结果
  reranked = rrf([keyword_search_results, vector_search_results])
  
  print(json.dumps(reranked, indent=4, ensure_ascii=False))
  
  
  from sentence_transformers import SentenceTransformer
  
  model_name = 'BAAI/bge-large-zh-v1.5' #中文
  # model_name = 'moka-ai/m3e-base'  # 中英双语，但效果一般
  # model_name = 'BAAI/bge-m3' # 多语言，但效果一般
  
  model = SentenceTransformer(model_name)
  
  query = "国际争端"
  # query = "global conflicts"
  
  documents = [
      "联合国就苏丹达尔富尔地区大规模暴力事件发出警告",
      "土耳其、芬兰、瑞典与北约代表将继续就瑞典“入约”问题进行谈判",
      "日本岐阜市陆上自卫队射击场内发生枪击事件 3人受伤",
      "国家游泳中心（水立方）：恢复游泳、嬉水乐园等水上项目运营",
      "我国首次在空间站开展舱外辐射生物学暴露实验",
  ]
  
  query_vec = model.encode(query)
  
  doc_vecs = [
      model.encode(doc)
      for doc in documents
  ]
  
  print("Cosine distance:")  # 越大越相似
  # print(cos_sim(query_vec, query_vec))
  for vec in doc_vecs:
      print(cos_sim(query_vec, vec))