Spaces:
Sleeping
Sleeping
| import os | |
| import torch | |
| from langchain import HuggingFacePipeline | |
| from transformers import AutoTokenizer | |
| import transformers | |
| from langchain import HuggingFaceHub | |
| def get_openai_chat_model(API_key): | |
| try: | |
| from langchain.llms import OpenAI | |
| except ImportError as err: | |
| raise "{}, unable to load openAI. Please install openai and add OPENAIAPI_KEY" | |
| os.environ["OPENAI_API_KEY"] = API_key | |
| llm = OpenAI() | |
| return llm | |
| def get_llama_model(temperature=0,api_key=None,max_tokens=2048): | |
| model_id = "meta-llama/Llama-2-7b-chat-hf" | |
| llm = None | |
| try: | |
| tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| pipeline = transformers.pipeline("text-generation", | |
| model=model_id, | |
| tokenizer=tokenizer, | |
| torch_dtype=torch.bfloat16, | |
| trust_remote_code=True, | |
| device_map="auto", | |
| max_length=1000, | |
| eos_token_id=tokenizer.eos_token_id | |
| ) | |
| llm = HuggingFacePipeline(pipeline = pipeline, model_kwargs = {'temperature':temperature}) | |
| except: | |
| raise "User not autorized to access the Model" | |
| return llm | |
| def get_model_from_hub(api_key,temperature=0.1,max_tokens=2048,model_id="meta-llama/Llama-2-7b-chat-hf"): | |
| llm = HuggingFaceHub(huggingfacehub_api_token=api_key, | |
| repo_id=model_id, | |
| model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens}) | |
| return llm |