受限环境下AWQ加载的本地开源LLM与Llama-Index集成方案问询

一、直接集成方案（无需自定义LLM类） #

Llama-Index的HuggingFaceLLM支持本地加载模型，你可以直接将AWQ加载好的模型和分词器传入该类完成集成：

from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.core.prompts.prompts import SimpleInputPrompt
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# 定义系统提示词与查询模板（可根据模型需求调整）
system_prompt = """你是专业助手，仅基于提供的上下文信息回答问题，不得编造内容。"""
query_wrapper_prompt = SimpleInputPrompt("{query_str}")

# 初始化HuggingFaceLLM
llm = HuggingFaceLLM(
    model=model,
    tokenizer=tokenizer,
    system_prompt=system_prompt,
    query_wrapper_prompt=query_wrapper_prompt,
    context_window=4096,  # 匹配你的模型上下文窗口大小
    max_new_tokens=512,  # 设置生成的最大token数
    generate_kwargs={"temperature": 0.1, "do_sample": False},
    device_map="auto"  # 硬件充足时可直接设为"cuda"
)

# 对接RAG流程
documents = SimpleDirectoryReader("本地文档路径").load_data()
index = VectorStoreIndex.from_documents(documents, llm=llm)
query_engine = index.as_query_engine()
response = query_engine.query("你的问题")
print(response)

二、自定义LLM类（可选，适配特殊需求） #

如果需要个性化生成逻辑（比如自定义输出格式、特殊token处理），可以继承Llama-Index的LLM基类实现自定义类：

from llama_index.core.llms import LLM, CompletionResponse
from llama_index.core.llms.callbacks import llm_completion_callback
import torch

class CustomAWQLLM(LLM):
    def __init__(self, model, tokenizer, context_window=4096, max_new_tokens=512):
        self.model = model
        self.tokenizer = tokenizer
        self.context_window = context_window
        self.max_new_tokens = max_new_tokens
        # 补全pad token（部分模型默认无该配置）
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

    @property
    def metadata(self):
        return {"context_window": self.context_window, "num_output": self.max_new_tokens}

    @llm_completion_callback()
    def complete(self, prompt: str, **kwargs) -> CompletionResponse:
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=self.context_window).to(self.model.device)
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=self.max_new_tokens,
                temperature=kwargs.get("temperature", 0.1),
                do_sample=kwargs.get("do_sample", False),
                pad_token_id=self.tokenizer.pad_token_id
            )
        response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)[len(prompt):].strip()
        return CompletionResponse(text=response_text)

# 实例化自定义LLM并对接RAG
custom_llm = CustomAWQLLM(model=model, tokenizer=tokenizer)
documents = SimpleDirectoryReader("本地文档路径").load_data()
index = VectorStoreIndex.from_documents(documents, llm=custom_llm)

三、最优方案推荐 #

优先选择**直接集成HuggingFaceLLM**的方案，Llama-Index已封装好通用生成逻辑，无需重复开发，能快速完成RAG流程对接。仅当需要特殊业务逻辑时，再考虑自定义LLM类。

内容的提问来源于stack exchange，提问作者Supersquirrel