update app.py

app.py

@@ -10,27 +10,21 @@ tokenizer = T5Tokenizer.from_pretrained("t5-base")
 model = T5ForConditionalGeneration.from_pretrained("t5-base")
 
 
-def parse_text(text):
-    sentences = sent_tokenize(text)
-    tokens = []
-    for sentence in sentences:
-        tokens.extend(sentence.split())
-    return ". ".join(tokens)
-
-
-def summarize(text):
-    input_ids = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
-    summary_ids = model.generate(input_ids=input_ids, max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
-    output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+def generate_response(text):
+    input_ids = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)
+    response_ids = model.generate(input_ids=input_ids, max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
+    output = tokenizer.decode(response_ids[0], skip_special_tokens=True)
     return output
 
 
-def format_messages_for_summary(messages):
-    summary_text_list = []
+def format_messages_for_display(messages):
+    formatted_text = []
     for message in messages:
         if message["role"] == "assistant":
-            summary_text_list.append(message["content"].lower())
-    return ' '.join(summary_text_list)
+            formatted_text.append(f"Assistant: {message['content']}")
+        else:
+            formatted_text.append(f"User: {message['content']}")
+    return "\n".join(formatted_text)
 
 
 def main():
@@ -51,14 +45,14 @@ def main():
             }
         ]
 
-        response = summarize(user_input)
+        response = generate_response(user_input)
 
         st.session_state['messages'].append({
             "role": "assistant",
             "content": response
         })
 
-        st.write(format_messages_for_summary(st.session_state['messages']))
+        st.write(format_messages_for_display(st.session_state['messages']))
 
 
 def save_session():