vaishakgkumar commited on
Commit
97822ab
1 Parent(s): 329f7fe

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +112 -0
app.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, MistralForCausalLM
2
+ import torch
3
+ import gradio as gr
4
+ import random
5
+ from textwrap import wrap
6
+ from transformers import AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, MistralForCausalLM
7
+ from peft import PeftModel, PeftConfig
8
+ import torch
9
+ import gradio as gr
10
+ import os
11
+
12
+ hf_token = os.environ.get('HUGGINGFACE_TOKEN')
13
+
14
+ # Functions to Wrap the Prompt Correctly
15
+ def wrap_text(text, width=90):
16
+ lines = text.split('\n')
17
+ wrapped_lines = [textwrap.fill(line, width=width) for line in lines]
18
+ wrapped_text = '\n'.join(wrapped_lines)
19
+ return wrapped_text
20
+ def multimodal_prompt(user_input, system_prompt="You are an expert medical analyst:"):
21
+
22
+ # Combine user input and system prompt
23
+ formatted_input = f"{user_input}{system_prompt}"
24
+
25
+ # Encode the input text
26
+ encodeds = tokenizer(formatted_input, return_tensors="pt", add_special_tokens=False)
27
+ model_inputs = encodeds.to(device)
28
+
29
+ # Generate a response using the model
30
+ output = model.generate(
31
+ **model_inputs,
32
+ max_length=max_length,
33
+ use_cache=True,
34
+ early_stopping=True,
35
+ bos_token_id=model.config.bos_token_id,
36
+ eos_token_id=model.config.eos_token_id,
37
+ pad_token_id=model.config.eos_token_id,
38
+ temperature=0.1,
39
+ do_sample=True
40
+ )
41
+
42
+ # Decode the response
43
+ response_text = tokenizer.decode(output[0], skip_special_tokens=True)
44
+
45
+ return response_text
46
+
47
+ # Define the device
48
+ device = "cuda" if torch.cuda.is_available() else "cpu"
49
+
50
+ # Use the base model's ID
51
+ base_model_id = "stabilityai/stablelm-3b-4e1t"
52
+ model_directory = "Tonic/stablemed"
53
+
54
+ # Instantiate the Tokenizer
55
+ tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t", token=hf_token, trust_remote_code=True, padding_side="left")
56
+ # tokenizer = AutoTokenizer.from_pretrained("Tonic/stablemed", trust_remote_code=True, padding_side="left")
57
+ tokenizer.pad_token = tokenizer.eos_token
58
+ tokenizer.padding_side = 'left'
59
+
60
+ # Load the PEFT model
61
+ peft_config = PeftConfig.from_pretrained("Tonic/stablemed", token=hf_token)
62
+ peft_model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t", token=hf_token, trust_remote_code=True)
63
+ peft_model = PeftModel.from_pretrained(peft_model, "Tonic/stablemed", token=hf_token)
64
+
65
+ class ChatBot:
66
+ def __init__(self):
67
+ self.history = []
68
+
69
+ def predict(self, user_input, system_prompt="You are an expert medical analyst:"):
70
+ # Combine user input and system prompt
71
+ formatted_input = f"{user_input}{system_prompt}"
72
+
73
+ # Encode user input
74
+ user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
75
+
76
+ # Concatenate the user input with chat history
77
+ if len(self.history) > 0:
78
+ chat_history_ids = torch.cat([self.history, user_input_ids], dim=-1)
79
+ else:
80
+ chat_history_ids = user_input_ids
81
+
82
+ # Generate a response using the PEFT model
83
+ response = peft_model.generate(input_ids=chat_history_ids, max_length=1200, pad_token_id=tokenizer.eos_token_id)
84
+
85
+ # Update chat history
86
+ self.history = chat_history_ids
87
+
88
+ # Decode and return the response
89
+ response_text = tokenizer.decode(response[0], skip_special_tokens=True)
90
+ return response_text
91
+
92
+ bot = ChatBot()
93
+
94
+ title = "👋🏻Welcome to Tonic's 😷StableMed⚕️ Chat🦟"
95
+ description = """
96
+ You can use this Space to test out the current model [StableMed](https://huggingface.co/Tonic/stablemed)
97
+ You can also use 😷StableMed⚕️ on your laptop & by cloning this space. 🧬🔬🔍 Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/StableMed_Chat?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3>
98
+ Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community on 👻Discord: [Discord](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to 🌟 [PolyGPT](https://github.com/tonic-ai/polygpt-alpha)
99
+ """
100
+ examples = [["What is the proper treatment for buccal herpes?", "Please provide information on the most effective antiviral medications and home remedies for treating buccal herpes."]]
101
+
102
+ iface = gr.Interface(
103
+ fn=bot.predict,
104
+ title=title,
105
+ description=description,
106
+ examples=examples,
107
+ inputs=["text", "text"], # Take user input and system prompt separately
108
+ outputs="text",
109
+ theme="ParityError/Anime"
110
+ )
111
+
112
+ iface.launch()