Upload app.py

app.py

@@ -7,11 +7,9 @@ import random
 import pandas as pd
 from sklearn.model_selection import train_test_split
 import time
-
 #from model import RNN_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
-
 ################################################################################
 import argparse
 import numpy as np
@@ -21,18 +19,13 @@ import copy
 from str2bool import str2bool
 from typing import Dict, Sequence
 from sentence_transformers import SentenceTransformer
-
 import torch
-import json
-
-import transformers
 from modeling_phi import PhiForCausalLM
 from tokenization_codegen import CodeGenTokenizer
+from transformers import PhiForCausalLM, AutoTokenizer, AutoModelForCausalLM
 ################################################################################
-
 parser = argparse.ArgumentParser()
 #############################################################################################################################
-
 parser.add_argument('--device_id', type=str, default="0")
 parser.add_argument('--model', type=str, default="microsoft/phi-2", help="") ## /phi-1.5
 parser.add_argument('--embedder', type=str, default="BAAI/bge-small-en-v1.5") ## /bge-small-en-v1.5  # bge-m3
@@ -51,7 +44,6 @@ if torch.cuda.is_available():
     print(f'################################################################# device: {device}#################################################################')
 else:
     device = "cpu"
-
 def get_model(base_model: str = "bigcode/starcoder",):
     tokenizer = CodeGenTokenizer.from_pretrained(base_model)
     tokenizer.pad_token_id = tokenizer.eos_token_id
@@ -59,14 +51,10 @@ def get_model(base_model: str = "bigcode/starcoder",):
 
     model = PhiForCausalLM.from_pretrained(
         base_model,
-        device_map="auto",
-    )
+        device_map="auto",)
     model.config.pad_token_id = tokenizer.pad_token_id
-
     model.eval()
-
     return tokenizer, model
-
 ################################################################################
 
 '''
@@ -87,14 +75,14 @@ The bot will respond based on the best possible answers to your messages.
 """
 # Create the gradio demo
 with gr.Blocks(css = """#col_container { margin-left: auto; margin-right: auto;} #chatbot {height: 520px; overflow: auto;}""") as demo:
-  gr.HTML('<h1 align="center">Medical Chatbot: ARIN 7102')
-  #gr.HTML('<h3 align="center">To know more about this project')
-  with gr.Accordion("Follow these Steps to use the Gradio WebUI", open=True):
-      gr.HTML(howto)
-  chatbot = gr.Chatbot()
-  msg = gr.Textbox()
-  clear = gr.ClearButton([msg, chatbot])
-
+    gr.HTML('<h1 align="center">Medical Chatbot: ARIN 7102')
+    #gr.HTML('<h3 align="center">To know more about this project')
+    with gr.Accordion("Follow these Steps to use the Gradio WebUI", open=True):
+        gr.HTML(howto)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    clear = gr.ClearButton([msg, chatbot])
+'''
   def respond(message, chat_history):
     # Create couple of if-else statements to capture/mimick peoples's Interaction
     embedder = SentenceTransformer(args.embedder, device=device)
@@ -104,10 +92,37 @@ with gr.Blocks(css = """#col_container { margin-left: auto; margin-right: auto;}
     chat_history.append((message, bot_message))
     time.sleep(2)
     return "", chat_history
-    
-  msg.submit(respond, [msg, chatbot], [msg, chatbot])
-  
+'''
+def respond(message, chat_history, base_model = "/home/henry/Desktop/ARIN7102/phi-2", device=device): # "meta-llama/Meta-Llama-3-70B"
+        # define the model and tokenizer.
+        # model = PhiForCausalLM.from_pretrained(base_model)
+        model = AutoModelForCausalLM.from_pretrained(base_model)
+        tokenizer = AutoTokenizer.from_pretrained(base_model)
+
+        # feel free to change the prompt to your liking.
+        #prompt = f"Patient: coercive spondylitis, pain in the lumbosacral area when turning over during sleep at night, no pain in any other part of the body.  
+        #/n Doctor: It shouldn't be a problem, but it's better to upload the images. /n Patient: {message} /n Doctor:"
+        output_termination = "\nOutput:"
+        prompt = f"Instruct: {message}{output_termination}"
+        # apply the tokenizer.
+        tokens = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
+        #tokens = tokens.to(device)
+        #eos_token_id = tokenizer.eos_token_id
+        # use the model to generate new tokens.
+        generated_output = model.generate(**tokens, use_cache=True, max_new_tokens=500, eos_token_id=50256, pad_token_id=50256)
+
+        # Find the position of "Output:" and extract the text after it
+        generated_text = tokenizer.batch_decode(generated_output)[0]
+        # Split the text at "Output:" and take the second part
+        split_text = generated_text.split("Output:", 1)
+        bot_message = split_text[1].strip() if len(split_text) > 1 else ""
+        bot_message = bot_message.replace("<|endoftext|>", "").strip()
+        chat_history.append((message, bot_message))
+        time.sleep(2)
+        return "", chat_history
+    #return bot_message
+
+msg.submit(respond, [msg, chatbot], [msg, chatbot])
 
 # Launch the demo
-demo.launch()
-
+demo.launch()