Update app.py

app.py

@@ -14,38 +14,6 @@ from huggingface_hub import login
 hf_token = os.environ.get('HUGGINGFACE_TOKEN')
 
 login(hf_token)
-# Functions to Wrap the Prompt Correctly
-def wrap_text(text, width=90):
-    lines = text.split('\n')
-    wrapped_lines = [textwrap.fill(line, width=width) for line in lines]
-    wrapped_text = '\n'.join(wrapped_lines)
-    return wrapped_text
-def multimodal_prompt(user_input, system_prompt="You are an expert medical analyst:"):
-
-    # Combine user input and system prompt
-    formatted_input = f"{user_input}{system_prompt}"
-
-    # Encode the input text
-    encodeds = tokenizer(formatted_input, return_tensors="pt", add_special_tokens=False)
-    model_inputs = encodeds.to(device)
-
-    # Generate a response using the model
-    output = model.generate(
-        **model_inputs,
-        max_length=max_length,
-        use_cache=True,
-        early_stopping=True,
-        bos_token_id=model.config.bos_token_id,
-        eos_token_id=model.config.eos_token_id,
-        pad_token_id=model.config.eos_token_id,
-        temperature=0.1,
-        do_sample=True
-    )
-
-    # Decode the response
-    response_text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return response_text
 
 # Define the device
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -67,35 +35,31 @@ peft_model = PeftModel.from_pretrained(peft_model, "vaishakgkumar/stablemedv1",
 class ChatBot:
     def __init__(self):
         self.history = []
-        
-    def predict(self, user_input, system_prompt="You are an expert analyst and provide assessment:"):
-        prompt = [{'role': 'user', 'content': user_input + "\n" + system_prompt + ":"}]
-        inputs = tokenizer.apply_chat_template(
-            prompt,
-            add_generation_prompt=True,
-            return_tensors='pt'
-        )
-
-        # Generate a response using the model
-        tokens = peft_model.generate(
-            inputs.to(model.device),
-            max_new_tokens=250,
-            temperature=0.8,
-            do_sample=False
-        )
-
-        # Decode the response
-        response_text = tokenizer.decode(tokens[0], skip_special_tokens=False)
-
-        # Free up memory
-        del tokens
-        torch.cuda.empty_cache()
 
-        return response_text
+    def predict(self, user_input, system_prompt="You are an expert medical analyst:"):
+        # Combine user input and system prompt
+        formatted_input = f"{system_prompt}{user_input}"
 
+        # Encode user input
+        user_input_ids = tokenizer.encode(formatted_input, return_tensors="pt")
 
-bot = ChatBot()
+        # Concatenate the user input with chat history
+        if len(self.history) > 0:
+            chat_history_ids = torch.cat([self.history, user_input_ids], dim=-1)
+        else:
+            chat_history_ids = user_input_ids
+
+        # Generate a response using the PEFT model
+        response = peft_model.generate(input_ids=chat_history_ids, max_length=1200, pad_token_id=tokenizer.eos_token_id)
+
+        # Update chat history
+        self.history = chat_history_ids
 
+        # Decode and return the response
+        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
+        return response_text
+
+bot = ChatBot()
 
 
 title = "StableDoc Chat"