SwatGarg commited on
Commit
5553e22
1 Parent(s): 2250554

Create generator.py

Browse files
Files changed (1) hide show
  1. generator.py +24 -0
generator.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from langchain_community.llms import HuggingFaceEndpoint
3
+
4
+
5
+ def load_llm(repo_id="mistralai/Mistral-7B-Instruct-v0.2"):
6
+ '''
7
+ Load the LLM from the HuggingFace model hub
8
+
9
+ Args:
10
+ repo_id (str): The HuggingFace model ID
11
+
12
+ Returns:
13
+ llm (HuggingFaceEndpoint): The LLM model
14
+ '''
15
+
16
+ repo_id = repo_id
17
+
18
+ llm = HuggingFaceEndpoint(
19
+ repo_id=repo_id, max_length=128, temperature=0.2, token=token)
20
+
21
+ return llm
22
+
23
+ def guardrails():
24
+ return None