Create custom-llama3.py

custom-llama3.py

@@ -0,0 +1,126 @@
+from typing import Any, Dict, List, Optional
+from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.language_models.llms import LLM
+import json
+import requests
+
+class CustomLlama3(LLM):
+    url: str = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
+    bearer_token: str
+    max_new_tokens: int = 512 
+    top_p: float = 0.9
+    temperature: float = 0.7
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        if stop is not None:
+            raise ValueError("stop kwargs are not permitted.")
+        json_body = {
+          "inputs": prompt,
+                  "parameters": {"max_new_tokens":self.max_new_tokens, "top_p":self.top_p, "temperature":self.temperature}
+          }
+        
+        data = json.dumps(json_body)
+        headers = {
+          'Content-Type': 'application/json',
+          'Authorization': f'Bearer {self.bearer_token}'
+        }
+        response = requests.request("POST", self.url, headers=headers, data=data)
+        response.raise_for_status()
+        try:
+          result=json.loads(response.content.decode("utf-8"))[0]['generated_text']
+          target_substring = '<|end_header_id|>'
+          substring_length = len(target_substring)
+          last_occurrence_index = result.rfind(target_substring)
+          if last_occurrence_index != -1:
+            return result[last_occurrence_index + substring_length:].strip()
+        except:
+          return response
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Return a dictionary of identifying parameters."""
+        return {"llmUrl": self.url}
+
+    @property
+    def _llm_type(self) -> str:
+        return "CustomLlama3"
+
+###########
+
+from langchain_core.prompts import PromptTemplate
+
+prompt = PromptTemplate(
+    template="""<|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a grader assessing relevance 
+    of a retrieved document to a user question. If the document contains keywords related to the user question, 
+    grade it as relevant. It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
+    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question. \n
+    Provide the binary score as a JSON with a single key 'score' and no premable or explanation.
+     <|eot_id|><|start_header_id|>user<|end_header_id|>
+    Here is the retrieved document: \n\n {document} \n\n
+    Here is the user question: {question} \n <|eot_id|><|start_header_id|>assistant<|end_header_id|>
+    """,
+    input_variables=["question", "document"],
+)
+
+##########
+
+from google.colab import userdata
+
+remote_llm = CustomLlama3(bearer_token = userdata.get('HF_TOKEN'))
+
+retrieval_grader = prompt | remote_llm 
+
+# Example usage
+question = "What are the health benefits of apples?"
+document = "Apples are rich in vitamins and fiber."
+
+result = retrieval_grader.invoke({
+    "question": question,
+    "document": document
+})
+
+print(result)
+
+###########
+
+from langchain_community.document_loaders import WebBaseLoader
+from langchain_community.vectorstores import Chroma
+from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+urls = [
+    "https://lilianweng.github.io/posts/2023-06-23-agent/",
+    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
+    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
+]
+
+docs = [WebBaseLoader(url).load() for url in urls]
+docs_list = [item for sublist in docs for item in sublist]
+
+text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+    chunk_size=250, chunk_overlap=0
+)
+doc_splits = text_splitter.split_documents(docs_list)
+
+embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
+
+# Add to vectorDB
+vectorstore = Chroma.from_documents(
+    documents=doc_splits,
+    collection_name="rag-chroma",
+    embedding=embedding_function,
+)
+retriever = vectorstore.as_retriever()
+
+question = "agent memory"
+docs = retriever.invoke(question)
+
+doc_txt = docs[3].page_content
+print(retrieval_grader.invoke({"question": question, "document": doc_txt}))
+