Update medchat.py

medchat.py

@@ -7,25 +7,23 @@ class MedChat:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.tokenizer = GPT2Tokenizer.from_pretrained(self.path)
         self.model = GPT2LMHeadModel.from_pretrained(self.path).to(self.device)
-    def prompt(self, input):
+
+
+    def forward(self, question):
         prompt_input = (
             "The conversation between human and AI assistant.\n"
-            f"[|Human|] {self.input}\n"
+            "[|Human|]"
             "[|AI|]"
         )
-        
-    def __call__(self, question):  
         sentence = prompt_input.format_map({'input': f"{question}"})
-        inputs = tokenizer(sentence, return_tensors="pt").to(self.device)
+        inputs = self.tokenizer(sentence, return_tensors="pt").to(self.device)
         with torch.no_grad():
             beam_output = self.model.generate(**inputs,
-                                        min_new_tokens=1, 
-                                        max_length=512,
-                                        num_beams=3,
-                                        repetition_penalty=1.2,
-                                        early_stopping=True,
-                                        eos_token_id=198 
-                                        )
-        return tokenizer.decode(beam_output[0], skip_special_tokens=True)
-
-
+                                              min_new_tokens=1,
+                                              max_length=512,
+                                              num_beams=3,
+                                              repetition_penalty=1.2,
+                                              early_stopping=True,
+                                              eos_token_id=198
+                                              )
+        return self.tokenizer.decode(beam_output[0], skip_special_tokens=True)