Update app.py

app.py

@@ -57,14 +57,34 @@ class Model:
         self.tokenizer = AutoTokenizer.from_pretrained('saadob12/t5_C2T_big')
         self.model = AutoModelForSeq2SeqLM.from_pretrained('saadob12/t5_C2T_big').to(self.device)
       elif mode.lower() == 'analytical':
-        self.tokenizer = AutoTokenizer.from_pretrained('saadob12/t5_C2T_autochart')
-        self.model = AutoModelForSeq2SeqLM.from_pretrained('saadob12/t5_C2T_autochart').to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained('saadob12/t5_autochart_2.0')
+        self.model = AutoModelForSeq2SeqLM.from_pretrained('saadob12/t5_autochart_2.0').to(self.device)
 
     def generate(self):
       tokens = self.tokenizer.encode(self.prefix + self.text,  truncation=self.truncation, padding=self.padding, return_tensors='pt').to(self.device)
       generated = self.model.generate(tokens, num_beams=4, max_length=256)
       tgt_text = self.tokenizer.decode(generated[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
       summary = str(tgt_text).strip('[]""')
+      if 'barchart' in summary:
+        summary.replace('barchart','statistic')
+      if 'bar graph' in summary:
+        summary.replace('bar graph','statistic')
+      if 'bar plot' in summary:
+        summary.replace('bar plot','statistic')
+      if 'scatter plot' in summary:
+        summary.replace('scatter plot','statistic')
+      if 'scatter graph' in summary:
+        summary.replace('scatter graph','statistic')
+      if 'scatterchart' in summary:
+        summary.replace('scatter chart','statistic')
+      if 'line plot' in summary:
+        summary.replace('line plot','statistic')
+      if 'line graph' in summary:
+        summary.replace('line graph','statistic')
+      if 'linechart' in summary:
+        summary.replace('linechart','statistic')
+      if 'graph' in summary:
+        summary.replace('graph','statistic'
       return summary
       
 st.title('Chart and Data Summarization')