update

app.py

@@ -1,13 +1,14 @@
 from langchain.chains import ConversationChain
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+from myLLM import AutoModelLanguageModel
 
-# Define the LangChain chat agent
 model_name = "microsoft/DialoGPT-medium"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+llm = AutoModelLanguageModel(model_name)
+
+agent = ConversationChain(llm=llm)
 
-agent = ConversationChain(model=model, tokenizer=tokenizer)
 
 # Define the Gradio interface
 def chatbot_interface(input_text):

myLLM.py

@@ -0,0 +1,18 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from langchain.chains import LanguageModel
+
+class AutoModelLanguageModel(LanguageModel):
+    def __init__(self, model_name_or_path):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+
+    def generate_prompt(self, input_text, history):
+        inputs = self.tokenizer.encode(input_text + self.tokenizer.eos_token, return_tensors="pt")
+        history = [self.tokenizer.encode(h + self.tokenizer.eos_token, return_tensors="pt") for h in history]
+        prompt = torch.cat(history + [inputs], dim=-1)
+        return prompt
+
+    def generate_response(self, prompt, max_length):
+        output = self.model.generate(prompt, max_length=max_length, pad_token_id=self.tokenizer.pad_token_id)
+        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return response