import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, AutoConfig
from peft import LoraConfig, PeftModel
import gradio as gr
from os.path import dirname

#model_path = "models/microsoft-phi2-custom" 
new_model = "garima-mahato/gm_oasst1_phi2_peft" #"./models/microsoft-phi2-custom/"  #"https://huggingface.co/spaces/garima-mahato/OAP2/tree/main/models/microsoft-phi-2-custom"
model_name = "microsoft/phi-2"
device_map = "auto" #{"": 0}

# Reload model in FP16 and merge it with LoRA weights
# base_model = AutoModelForCausalLM.from_pretrained(
#     model_name,
#     low_cpu_mem_usage=True,
#     return_dict=True,
#     torch_dtype=torch.float16
#     #device_map=device_map
# )
phi_model = AutoModelForCausalLM.from_pretrained(new_model, trust_remote_code=True) # PeftModel.from_pretrained(base_model, new_model)
phi_tokenizer = AutoTokenizer.from_pretrained(new_model, trust_remote_code=True)
# AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# phi_tokenizer.pad_token = phi_tokenizer.eos_token
# phi_tokenizer.padding_side = "right"

def generate_answer(prompt):
    gen = pipeline('text-generation', model=phi_model, tokenizer=phi_tokenizer)
    result = gen(prompt)
    
    return result[0]['generated_text'].replace(prompt, '')

title = "Phi2-enabled Virtual Assistant"
description = "Ask Anything."
examples = ["What is monospony?","What is monospony in economics?"]

question = gr.TextArea(label="Do you want to ask any question?")
answer = gr.TextArea(label="Your answer")

demo = gr.Interface(
    generate_answer, 
    inputs = question, 
    outputs = answer,
    title = title,
    description = description,
    examples = examples
)
demo.launch()