import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import streamlit as st
import accelerate
import re

def formatting_func(document):
    instruction = "You are a model designed to rephrase medical summaries for a general audience. Please summarize the following article in such a way a normal person could understand it, while also ensuring the same factual accuracy. Replace any technical terms with their equivalents in ordinary language, and be concise (< 100 words) and approachable.\n---------\n" 
    text = f"### {instruction} \n### Conclusion: {document} \n### Summary: "
    return text

def generate(text,max_new_token):
    ft_model = AutoModelForCausalLM.from_pretrained("BiswajitPadhi99/mistral-7b-finetuned-medical-summarizer-old", 
                                                device_map="cuda", load_in_4bit=True)
    eval_tokenizer = AutoTokenizer.from_pretrained("BiswajitPadhi99/mistral-7b-finetuned-medical-summarizer-old", add_bos_token=True,
                                               device_map="cuda")
    ft_model.eval()
    with torch.no_grad():
        eval_prompt = formatting_func(text)
        model_input = eval_tokenizer(eval_prompt, return_tensors="pt").to("cuda")
        response = eval_tokenizer.decode(ft_model.generate(**model_input, max_new_tokens=max_new_token)[0], skip_special_tokens=True)

        if(eval_prompt in response):
            response = response.replace(eval_prompt, '')
        response = re.sub(r'#+', '', response)
        return response

def main():
    st.title('Medical Document Summarization')

    col1, col2 = st.columns(2)

    with col1:
        user_input = st.text_area("Enter your text here:", height=300)
        max_new_token = st.number_input('Max new tokens:', value=200)
        submit_button = st.button("Summarize")

    with col2:
        if submit_button:
            st.write("Model Response:")
            output = generate(user_input,max_new_token)
            print(output)
            st.markdown(output)

st.set_page_config(layout="wide")
if __name__=="__main__":
    main()