update

README.md

@@ -53,4 +53,4 @@ Each member should create a branch with their own name and commit to that branch
 
 ### Bug Fixes and Questions
 
-For bug fixes or questions, either open an issue or create a branch prefixed with `bug` in name.
+For bug fixes or questions, either open an issue or create a branch prefixed with `bug` in name.

pages/summarization_example.py

@@ -5,10 +5,14 @@ import random
 from utils.model import Model
 from utils.data import dataset
 
-__default_model_name__ = "lmsys/vicuna-7b-v1.5"
-model = Model(__default_model_name__)
 load_dotenv()
 
+__model_list__ = [
+    "lmsys/vicuna-7b-v1.5",
+    "tiiuae/falcon-7b-instruct"
+]
+model = {model_name: Model(model_name) for model_name in __model_list__}
+
 random_label = '🔀 Random dialogue from dataset'
 examples = {
     "example 1": """Boston's injury reporting for Kristaps Porziņģis has been fairly coy. He missed Game 3, but his coach told reporters just before Game 4 that was technically available, but with a catch.
@@ -65,7 +69,7 @@ summarization: """, label='Input Prompting Template', lines=8, placeholder='Inpu
         output = gr.Markdown()
 
         example_dropdown.change(update_input, inputs=[example_dropdown], outputs=[input_text])
-        submit_button.click(process_input, inputs=[input_text, model_dropdown, Template_text], outputs=[output])
+        submit_button.click(process_input, inputs=[input_text, model[model_dropdown], Template_text], outputs=[output])
 
     return demo
 

utils/model.py

@@ -1,20 +1,37 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
 import transformers
 import torch
 
-
 class Model():
-    def __init__(self, model="lmsys/vicuna-7b-v1.5") -> None:
-        self.tokenizer = AutoTokenizer.from_pretrained(model)
+    number_of_models = 0
+
+    def __init__(self, model_name="lmsys/vicuna-7b-v1.5") -> None:
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.name = model_name
         self.pipeline = transformers.pipeline(
             "text-generation",
-            model=model,
+            model=model_name,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
             device_map="auto",
         )
 
+        self.update()
+
+    @classmethod
+    def update(cls):
+        cls.number_of_models += 1
+
+    def return_mode_name(self):
+        return self.name
+    
+    def return_tokenizer(self):
+        return self.tokenizer
+    
+    def return_model(self):
+        return self.pipeline
+
     def gen(self, content, temp=0.1, max_length=500):
         sequences = self.pipeline(
             content,
@@ -26,4 +43,4 @@ class Model():
             return_full_text=False
         )
         
-        return sequences[-1]['generated_text'] #'\n'.join([seq['generated_text'] for seq in sequences])
+        return sequences[-1]['generated_text']