import random
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

tokenizer = AutoTokenizer.from_pretrained("docto/Docto-Bot")
model = AutoModelForCausalLM.from_pretrained("docto/Docto-Bot")
special_token = '<|endoftext|>'


def get_reply(userinput):
        prompt_text = f'Question: {userinput}\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.7,
            top_k = 20,
            top_p = 0.9,
            repetition_penalty = 1,
            do_sample = True,
            num_return_sequences = 1
        )

        # 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"
            
iface = gr.Interface(fn=get_reply, inputs=["text"], outputs=["textbox"]).launch()