Update app.py

app.py

@@ -1,4 +1,4 @@
-import gradio as gr           
+import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
@@ -10,48 +10,41 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=bnb_config)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id,  # Mistral, same as before
-    quantization_config=bnb_config,  # Same quantization config as before
-    device_map="auto",
-    trust_remote_code=True,
-    use_auth_token=api_token
-)
+# Load the fine-tuned model "Tonic/mistralmed"
+model = AutoModelForCausalLM.from_pretrained("Tonic/mistralmed", quantization_config=bnb_config)
+
+tokenizer = AutoTokenizer.from_pretrained("Tonic/mistralmed", trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = 'left'
+
+class ChatBot:
+    def __init__(self):
+        self.history = []
+
+    def predict(self, input):
+        new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")
+        flat_history = [item for sublist in self.history for item in sublist]
+        flat_history_tensor = torch.tensor(flat_history).unsqueeze(dim=0)
+        bot_input_ids = torch.cat([flat_history_tensor, new_user_input_ids], dim=-1) if self.history else new_user_input_ids
+        chat_history_ids = model.generate(bot_input_ids, max_length=2000, pad_token_id=tokenizer.eos_token_id)
+        self.history.append(chat_history_ids[:, bot_input_ids.shape[-1]:].tolist()[0])
+        response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+        return response
 
-tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_tokentokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")    
-tokenizer.padding_side = 'left'      
-model = PeftModel.from_pretrained(base_model, "Tonic/mistralmed")
-    
-class ChatBot:      
-    def __init__(self):      
-        self.history = []      
-    
-    def predict(self, input):            
-        new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")            
-        flat_history = [item for sublist in self.history for item in sublist]       
-        flat_history_tensor = torch.tensor(flat_history).unsqueeze(dim=0)  # convert list to 2-D tensor    
-        bot_input_ids = torch.cat([flat_history_tensor, new_user_input_ids], dim=-1) if self.history else new_user_input_ids            
-        chat_history_ids = model.generate(bot_input_ids, max_length=2000, pad_token_id=tokenizer.eos_token_id)            
-        self.history.append(chat_history_ids[:, bot_input_ids.shape[-1]:].tolist()[0])            
-        response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)     
-        return response      
-  
-bot = ChatBot() 
-  
-title = "👋🏻Welcome to Tonic's EZ Chat🚀"    
-description = "You can use this Space to test out the current model (MistralMed) or duplicate this Space and use it for any other model on 🤗HuggingFace. Join me on [Discord](https://discord.gg/fpEPNZGsbt) to build together."    
-examples = [["What is the boiling point of nitrogen?"]]    
-  
-iface = gr.Interface(    
-    fn=bot.predict,    
-    title=title,    
-    description=description,    
-    examples=examples,    
-    inputs="text",    
-    outputs="text", 
+bot = ChatBot()
+
+title = "👋🏻Welcome to Tonic's EZ Chat🚀"
+description = "You can use this Space to test out the current model (MistralMed) or duplicate this Space and use it for any other model on 🤗HuggingFace. Join me on [Discord](https://discord.gg/fpEPNZGsbt) to build together."
+examples = [["What is the boiling point of nitrogen"]]
+
+iface = gr.Interface(
+    fn=bot.predict,
+    title=title,
+    description=description,
+    examples=examples,
+    inputs="text",
+    outputs="text",
     theme="ParityError/Anime"
-)    
-  
-iface.launch()  
+)
+
+iface.launch()