Update app.py

app.py

@@ -80,19 +80,7 @@ class TeapotAI:
         if self.settings.verbose:
             print(f"Loading Model: {self.model} Revision: {self.model_revision or 'Latest'}")
         
-        # self.generator = pipeline("text2text-generation", model=self.model, revision=self.model_revision) if model_revision else pipeline("text2text-generation", model=self.model)
-
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model)
-        model = AutoModelForSeq2SeqLM.from_pretrained(self.model)
-        model.eval()
-
-         # Quantization settings 
-        quantization_dtype = torch.qint8  # or torch.float16
-        quantization_config = torch.quantization.get_default_qconfig('fbgemm')  # or 'onednn'
-
-        self.quantized_model = torch.quantization.quantize_dynamic(
-            model, {torch.nn.Linear}, dtype=quantization_dtype
-        )
+        self.generator = pipeline("text2text-generation", model=self.model, revision=self.model_revision) if model_revision else pipeline("text2text-generation", model=self.model)
 
         self.documents = documents
         
@@ -154,14 +142,8 @@ class TeapotAI:
             str: The generated output from the model.
         """
         
-        inputs = self.tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-
-          
-        with torch.no_grad():
-            outputs = self.quantized_model.generate(inputs["input_ids"], max_length=512)
         
-
-        result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        result = self.generator(input_text)[0].get("generated_text")
         
 
         if self.settings.log_level == "debug":