Update app.py

app.py

@@ -2,31 +2,43 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 model_id = "Mohammed-Altaf/medical_chatbot-8bit"
-model = AutoModelForCausalLM.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id,ignore_mismatched_sizes=True)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
+def get_clean_response(response):
+    if type(response) == list:
+        response = response[0].split("\n")
+    else:
+        response = response.split("\n")
+        
+    ans = ''
+    cnt = 0 # to verify if we have seen Human before 
+    for answer in response:
+        if answer.startswith("[|Human|]"): cnt += 1
+            
+        elif answer.startswith('[|AI|]'):
+            answer = answer.split(' ')
+            ans += ' '.join(char for char in answer[1:])
+            ans += '\n'
+        
+        elif cnt:
+            ans += answer + '\n'
+    return ans
+
+
 def generate_text(input_text):
     input_ids = tokenizer.encode(input_text, return_tensors="pt")
 
     output = model.generate(
-        input_ids,
-        max_length=200,
+        **input_ids,
+        max_length=100,
         )
 
     output_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    print(output_text)
 
-    # Remove Prompt Echo from Generated Text
-    cleaned_output_text = output_text.replace(input_text, "")
-    return cleaned_output_text
+    return get_clean_response(output_text)
 
 
-text_generation_interface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.inputs.Textbox(label="Input Text"),
-    ],
-    outputs=gr.inputs.Textbox(label="Generated Text"),
-    title="Medical ChatBot",
-).launch()
+iface = gr.Interface(fn = generate_text, inputs = 'text', outputs = ['text'], title ='Medical ChatBot')
+iface.launch()