Update app.py

app.py

@@ -1,12 +1,20 @@
 import gradio as gr
-from transformers import T5TokenizerFast
-from model import PoemSummaryModel
-
+import torch
+from transformers import T5ForConditionalGeneration, T5TokenizerFast
 
 tokenizer = T5TokenizerFast.from_pretrained("t5-base")
-best_model = PoemSummaryModel.load_from_checkpoint("best-checkpoint.ckpt")
-best_model.freeze()
 
+# Define the quantized model architecture
+quantized_model = T5ForConditionalGeneration.from_pretrained("t5-base")
+
+# Load the state dictionary
+state_dict = torch.load("quantized_model.pt")
+
+# Filter out keys that are not present in the quantized model
+filtered_state_dict = {k: v for k, v in state_dict.items() if k in quantized_model.state_dict()}
+
+# Load the filtered state dictionary into the quantized model
+quantized_model.load_state_dict(filtered_state_dict, strict=False)
 
 def encode_text(text):
     encoding = tokenizer.encode_plus(
@@ -21,7 +29,7 @@ def encode_text(text):
 
 def generate_summary(input_ids, attention_mask, model):
     model = model.to(input_ids.device)
-    generated_ids = model.model.generate(
+    generated_ids = model.generate(
         input_ids=input_ids,
         attention_mask=attention_mask,
         max_length=150,
@@ -39,7 +47,7 @@ def decode_summary(generated_ids):
 
 def summarize(text):
     input_ids, attention_mask = encode_text(text)
-    generated_ids = generate_summary(input_ids, attention_mask, best_model)
+    generated_ids = generate_summary(input_ids, attention_mask, quantized_model)
     summary = decode_summary(generated_ids)
     return summary
 
@@ -53,4 +61,4 @@ gr.Interface(
     outputs=output_text,
     title="Poem Pulse",
     description="Enter a Poem and get its Jist."
-).launch()
+).launch()