test model

pages/summarization_example.py

@@ -2,6 +2,11 @@ from dotenv import load_dotenv
 from transformers import pipeline
 import gradio as gr
 
+from utils.model import Model
+
+
+model = Model()
+
 load_dotenv()
 
 examples = {
@@ -26,11 +31,11 @@ def generate_answer(sources, model_name):
 summarization: """
     content = meta_prompt.format(sources=sources)
 
-    pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True, max_length=500)
-    messages = [{"role": "user", "content": content}]
-    answer = pipe(messages)
+    # pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True, max_length=500)
+    # messages = [{"role": "user", "content": content}]
+    answer = model.gen(content)
 
-    return answer[0]['generated_text']
+    return answer
 
 def process_input(input_text, model_selection):
     if input_text:

utils/__init__.py

@@ -1,4 +1,4 @@
 # This is the __init__.py file for the utils package
 # You can add any initialization code or import statements here
 
-__all__ = ['multiple_stream']
+__all__ = ['multiple_stream', 'model']

utils/model.py

@@ -0,0 +1,29 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import transformers
+import torch
+
+
+class Model():
+    def __init__(self, model="tiiuae/falcon-7b-instruct") -> None:
+        pass
+        self.tokenizer = AutoTokenizer.from_pretrained(model)
+        self.pipeline = transformers.pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=self.tokenizer,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+            device_map="auto",
+        )
+
+    def gen(self, prompt, temp=0.0, max_length=200):
+        sequences = self.pipeline(
+            "Girafatron is obsessed with giraffes, the most glorious animal on the face of this Earth. Giraftron believes all other animals are irrelevant when compared to the glorious majesty of the giraffe.\nDaniel: Hello, Girafatron!\nGirafatron:",
+            max_length=max_length,
+            do_sample=True,
+            temperature=temp,
+            num_return_sequences=1,
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        
+        return '\n'.join([seq['generated_text'] for seq in sequences])