Update model_list.py

model_list.py

@@ -16,6 +16,8 @@ class ModelList:
         self.table_header = '''
             <tr>
                 <td width="20%">Model Name</td>
+                <td width="10%">Type</td>
+                <td width="10%">Base Model</td>
                 <td width="10%">Paper</td>
                 <td width="10%">Code on Github</td>
                 <td width="10%">Weights on 🤗</td>
@@ -35,9 +37,15 @@ class ModelList:
                 row.hub, str) else ''
             other_model = f'<a href="{row.other}" target="_blank">Other Weights</a>' if isinstance(
                 row.other, str) else ''
+            data_type = f'<a href="{row.type}" target="_blank">Paper</a>' if isinstance(
+                row.paper, str) else ''
+            base_model = f'<a href="{row.base_model}" target="_blank">Paper</a>' if isinstance(
+                row.paper, str) else ''
             row = f'''
                 <tr>
                     <td>{row.name}</td>
+                    <td>{data_type}</td>
+                    <td>{base_model}</td>
                     <td>{paper}</td>
                     <td>{github}</td>
                     <td>{hf_model}</td>
@@ -49,7 +57,8 @@ class ModelList:
     def render(self, search_query: str, 
             case_sensitive: bool,
             filter_names: list[str],
-            data_types: list[str]) -> tuple[int, str]:
+            data_types: list[str],
+            model_types: list[str]) -> tuple[int, str]:
         df = self.table
         if search_query:
             if case_sensitive:
@@ -59,7 +68,7 @@ class ModelList:
         has_paper = 'Paper' in filter_names
         has_github = 'Code' in filter_names
         has_model = 'Model Weights' in filter_names
-        df = self.filter_table(df, has_paper, has_github, has_model, data_types)
+        df = self.filter_table(df, has_paper, has_github, has_model, data_types, model_types)
         return len(df), self.to_html(df, self.table_header)
 
     @staticmethod
@@ -71,7 +80,8 @@ class ModelList:
             df = df[~df.github.isna()]
         if has_model:
             df = df[~df.hub.isna() | ~df.other.isna()]
-        #df = df[df.type.isin(set(data_types))]
+        df = df[df.type.isin(set(data_types))]
+        df = df[df.type.isin(set(model_types))]
         return df
 
     @staticmethod