add models

models/distilbart_cnn_12_6.py

@@ -0,0 +1,47 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# Load the DistilBART-CNN-12-6 model
+# loading the model outside of the function makes it faster
+SUMMARIZATION_MODEL = "sshleifer/distilbart-cnn-12-6"
+tokenizer   = AutoTokenizer.from_pretrained(SUMMARIZATION_MODEL)
+model       = AutoModelForSeq2SeqLM.from_pretrained(SUMMARIZATION_MODEL, device_map="cuda:0")
+
+def summarize(text, max_len=20):
+    """
+    Summarizes the given text using the DistilBART-CNN-12-6 model.
+
+    Args:
+        text (str): The text to be summarized.
+        max_length (int, optional): The maximum length of the summary. Defaults to 20.
+
+    Returns:
+        str: The summarized text.
+    """
+    
+    inputs = tokenizer(text, 
+                       return_tensors="pt",
+                       max_length=max_len,
+                       truncation=True,
+    ).input_ids
+    
+    # Move the inputs tensor to the same device as the model tensor
+    inputs = inputs.cuda()
+    
+    outputs = model.generate(inputs, 
+                            max_new_tokens=100, 
+                            num_beams=8, 
+                            length_penalty=0.2, 
+                            early_stopping=False
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def summarizePipeline(text):
+    from transformers import pipeline
+    
+    pipe = pipeline(
+        "summarization",
+        model=model,
+        tokenizer=tokenizer,
+    )
+    
+    return pipe(text)[0]["summary_text"]

models/flan_t5_xl.py

@@ -0,0 +1,16 @@
+import os, torch, accelerate
+from langchain.llms import HuggingFacePipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, AutoModelForCausalLM
+
+model_id = 'google/flan-t5-large'
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id, load_in_8bit=True)
+
+pipe = pipeline(
+    "text2text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_length=512,
+)
+
+local_llm = HuggingFacePipeline(Pipeline=pipe)

models/t5_small_medium_title_generation.py

@@ -0,0 +1,30 @@
+import nltk
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+def t5model(prompt: str) -> str:
+    tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-small-medium-title-generation")
+    model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation", device_map="cuda:0", torch_dtype=torch.float16)
+    inputs = tokenizer(
+        ["summarize:" + prompt],
+        return_tensors="pt",
+        max_length=1024,
+        truncation=True
+    )
+    
+    # Move the inputs tensor to the same device as the model tensor
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    
+    outputs = model.generate(
+        **inputs,
+        num_beams=8,
+        do_sample=True,
+        min_length=8,
+        max_length=15
+    )
+    
+    decoded_output = tokenizer.batch_decode(
+        outputs, skip_special_tokens=True
+    )[0]
+    
+    return decoded_output

requirements.txt

@@ -6,3 +6,4 @@ accelerate
 langchain
 yt-dlp
 rich
+gradio