LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-12-02 10:03:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #425 from Archit-Kohli/patch-2

Browse Source 
Update huggingface.py
This commit is contained in:
Alex
2023-10-05 11:03:27 +01:00
committed by GitHub
parent 6c3ed5e533 7ce1dc9069
commit 4b629d20cf
1 changed files with 18 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
application/llm/huggingface.py
View File

@@ -2,13 +2,26 @@ from application.llm.base import BaseLLM
class HuggingFaceLLM(BaseLLM):
def __init__(self, api_key, llm_name='Arc53/DocsGPT-7B'):
def __init__(self, api_key, llm_name='Arc53/DocsGPT-7B',q=False):
global hf
from langchain.llms import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
tokenizer = AutoTokenizer.from_pretrained(llm_name)
model = AutoModelForCausalLM.from_pretrained(llm_name)
if q:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
tokenizer = AutoTokenizer.from_pretrained(llm_name)
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
model = AutoModelForCausalLM.from_pretrained(llm_name,quantization_config=bnb_config)
else:
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
tokenizer = AutoTokenizer.from_pretrained(llm_name)
model = AutoModelForCausalLM.from_pretrained(llm_name)
pipe = pipeline(
"text-generation", model=model,
tokenizer=tokenizer, max_new_tokens=2000,