code: import os
  from fedml.serving import FedMLPredictor
  from fedml.serving import FedMLInferenceRunner
  from langchain import PromptTemplate, LLMChain
  from langchain.llms import HuggingFacePipeline
  import torch
  from transformers import (
      AutoConfig,
      AutoModelForCausalLM,
      AutoTokenizer,
      TextGenerationPipeline,
  )
  
  class Chatbot(FedMLPredictor):                # Inherit FedMLClientPredictor
      def __init__(self):
          super().__init__()
          PROMPT_FOR_GENERATION_FORMAT = f""""Below is an instruction that describes a task. Write a response that appropriately completes the request."
  
          ### Instruction:
          {{instruction}}
  
          ### Response:
          """
  
          prompt = PromptTemplate(
              input_variables=["instruction"],
              template=PROMPT_FOR_GENERATION_FORMAT
          )
  
          config = AutoConfig.from_pretrained("EleutherAI/pythia-70m")
          model = AutoModelForCausalLM.from_pretrained(
              "EleutherAI/pythia-70m",
              torch_dtype=torch.float32,      # float 16 not supported on CPU
              trust_remote_code=True,
              device_map="auto"
          )
          tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m", device_map="auto")
  
          hf_pipeline = HuggingFacePipeline(
              pipeline=TextGenerationPipeline(
                  model=model,
                  tokenizer=tokenizer,
                  return_full_text=True,
                  task="text-generation",
                  do_sample=True,
                  max_new_tokens=256,
                  top_p=0.92,
                  top_k=0
              )
          )
          self.chatbot = LLMChain(llm=hf_pipeline, prompt=prompt, verbose=True)
      
      def predict(self, request:dict):
          input_dict = request
          question: str = input_dict.get("text", "").strip()
  
          if len(question) == 0:
              response_text = "<received empty input; no response generated.>"
          else:
              response_text = self.chatbot.predict(instruction=question)
  
          return {"generated_text": str(response_text)}
  
  if __name__ == "__main__":
      print("Program starts...")
  
      # Parse arguments
      import argparse
      parser = argparse.ArgumentParser()
      parser.add_argument("--batch_size", type=int, default=50051, help="port number")
      args = parser.parse_args()
      print(f"args.batch_size: {args.batch_size}")
  
      # Parse environment variables
      local_rank = int(os.environ.get("LOCAL_RANK", 100))
      print(f"local rank: {local_rank}")
  
      chatbot = Chatbot()
      fedml_inference_runner = FedMLInferenceRunner(chatbot)
      fedml_inference_runner.run()
apis: langchain.LLMChain langchain.PromptTemplate
extract_api: [((2184, 2209), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (2207, 2209), False, 'import argparse\n'), ((2559, 2588), 'fedml.serving.FedMLInferenceRunner', 'FedMLInferenceRunner', (['chatbot'], {}), '(chatbot)\n', (2579, 2588), False, 'from fedml.serving import FedMLInferenceRunner\n'), ((706, 797), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['instruction']", 'template': 'PROMPT_FOR_GENERATION_FORMAT'}), "(input_variables=['instruction'], template=\n    PROMPT_FOR_GENERATION_FORMAT)\n", (720, 797), False, 'from langchain import PromptTemplate, LLMChain\n'), ((845, 896), 'transformers.AutoConfig.from_pretrained', 'AutoConfig.from_pretrained', (['"""EleutherAI/pythia-70m"""'], {}), "('EleutherAI/pythia-70m')\n", (871, 896), False, 'from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline\n'), ((913, 1049), 'transformers.AutoModelForCausalLM.from_pretrained', 'AutoModelForCausalLM.from_pretrained', (['"""EleutherAI/pythia-70m"""'], {'torch_dtype': 'torch.float32', 'trust_remote_code': '(True)', 'device_map': '"""auto"""'}), "('EleutherAI/pythia-70m', torch_dtype=\n    torch.float32, trust_remote_code=True, device_map='auto')\n", (949, 1049), False, 'from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline\n'), ((1160, 1233), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['"""EleutherAI/pythia-70m"""'], {'device_map': '"""auto"""'}), "('EleutherAI/pythia-70m', device_map='auto')\n", (1189, 1233), False, 'from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline\n'), ((1635, 1689), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'hf_pipeline', 'prompt': 'prompt', 'verbose': '(True)'}), '(llm=hf_pipeline, prompt=prompt, verbose=True)\n', (1643, 1689), False, 'from langchain import PromptTemplate, LLMChain\n'), ((2431, 2464), 'os.environ.get', 'os.environ.get', (['"""LOCAL_RANK"""', '(100)'], {}), "('LOCAL_RANK', 100)\n", (2445, 2464), False, 'import os\n'), ((1299, 1469), 'transformers.TextGenerationPipeline', 'TextGenerationPipeline', ([], {'model': 'model', 'tokenizer': 'tokenizer', 'return_full_text': '(True)', 'task': '"""text-generation"""', 'do_sample': '(True)', 'max_new_tokens': '(256)', 'top_p': '(0.92)', 'top_k': '(0)'}), "(model=model, tokenizer=tokenizer, return_full_text=\n    True, task='text-generation', do_sample=True, max_new_tokens=256, top_p\n    =0.92, top_k=0)\n", (1321, 1469), False, 'from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline\n')]