Update app.py

app.py

@@ -3,4 +3,31 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("docto/Docto-Bot")
 model = AutoModelForCausalLM.from_pretrained("docto/Docto-Bot")
-special_token = '<|endoftext|>'
+special_token = '<|endoftext|>'
+
+def get_reply(user_input):
+        prompt_text = f'Question: {user_input}\nAnswer:'
+        encoded_prompt = tokenizer.encode(prompt_text,
+                                  add_special_tokens = False,
+                                  return_tensors = 'pt')
+
+        output_sequences = model.generate(
+            input_ids = encoded_prompt,
+            max_length = 500,
+            temperature = 0.9,
+            top_k = 20,
+            top_p = 0.9,
+            repetition_penalty = 1,
+            do_sample = True,
+            num_return_sequences = 4
+        )
+
+        # result = tokenizer.decode(random.choice(output_sequences))
+        # result = result[result.index("Answer: "):result.index(special_token)]
+        try:
+            result = tokenizer.decode(random.choice(output_sequences))
+            result = result[result.index("Answer: "):result.index(special_token)]
+            return(result[8:])
+        except:
+            return "Sorry! I don\'t Know"
+