z

app.py

@@ -1,4 +1,4 @@
-from transformers import GPT2LMHeadModel, AutoTokenizer, pipeline
+from transformers import GPT2LMHeadModel, AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import gradio as gr
 
@@ -20,66 +20,86 @@ def generate_text_stream(model, tokenizer, prompt, max_new_tokens, temperature):
     inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     generated_tokens = 0
     output_text = ""
-
     while generated_tokens < max_new_tokens:
-        # Generate one token at a time
         outputs = model.generate(
             inputs,
-            max_new_tokens=1,  # Generate one token at a time
+            max_new_tokens=1,
             do_sample=True,
             top_p=0.95,
             top_k=50,
             temperature=temperature,
             pad_token_id=tokenizer.eos_token_id,
         )
-        # Decode the new token
         new_token = tokenizer.decode(outputs[0, -1], skip_special_tokens=True)
         output_text += new_token
         generated_tokens += 1
-
-        # Yield the updated output text
         yield output_text
-
-        # Update inputs for the next iteration
         inputs = outputs
 
-# Function to summarize text
-def summarize(text, summarizer, max_length, min_length):
-    summary = summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)
-    return summary[0]['summary_text']
+# Load BART model and tokenizer for summarization
+summarization_model_name = "facebook/bart-large-cnn"
+summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(summarization_model_name).to(torch.device("cuda" if torch.cuda.is_available() else "cpu")) # Use GPU if available
+summarizer_tokenizer = AutoTokenizer.from_pretrained(summarization_model_name)
+
+# Function to generate summary with manual streaming
+def generate_summary_stream(model, tokenizer, text, max_length, min_length):
+    inputs = tokenizer(text, return_tensors="pt", max_length=1024, truncation=True).to(model.device)
+    generated_tokens = 0
+    summary_text = ""
+    decoder_start_token_id = model.config.decoder_start_token_id
+    decoder_input_ids = torch.tensor([[decoder_start_token_id]], device=model.device)
+    past_key_values = None
+    while True :
+        outputs = model.generate(
+            inputs,
+            decoder_input_ids=decoder_input_ids,
+            max_new_tokens=1,
+            min_length=min_length,
+            do_sample=False,
+            past_key_values=past_key_values,
+            output_hidden_states=True,
+            return_dict_in_generate=True
+        )
+        next_token_id = outputs.sequences[0][-1]
+        if next_token_id == tokenizer.eos_token_id or len(decoder_input_ids[0]) >= max_length:
+            break
+        next_token = tokenizer.decode(next_token_id, skip_special_tokens=True)
+        summary_text += next_token
+        yield summary_text
+        decoder_input_ids = torch.cat([decoder_input_ids, torch.tensor([[next_token_id]], device=model.device)], dim=-1)
+        past_key_values = outputs.past_key_values
 
 def reply(prompt):
     if len(prompt) == 0:
         prompt = "class holocaust"
-    # Stream output from the first model (deleuze)
-    output1 = ""
-    for text in generate_text_stream(deleuze, tokenizer1, prompt, max_new_tokens=500, temperature=0.9):
-        output1 = text
-        yield output1
 
-    # Stream output from the second model (scum)
-    output2 = ""
-    for text in generate_text_stream(scum, tokenizer2, output1, max_new_tokens=200, temperature=1.7):
-        output2 = text
-        yield output2
+    # --- Phase 1: Generate and Stream Combined Output ---
+    combined_output = ""
+    scum_output = ""
+    gospel_output = ""
 
-    # Stream output from the third model (gospel)
-    output3 = ""
-    for text in generate_text_stream(gospel, tokenizer3, prompt, max_new_tokens=200, temperature=1.0):
-        output3 = text
-        yield output3
+    # Stream deleuze output
+    for text in generate_text_stream(deleuze, tokenizer1, prompt, max_new_tokens=500, temperature=0.9):
+      combined_output = text
+      yield combined_output
 
-    # Combine outputs for summarization
-    final_output = output2 + " " + output3
 
-    # Initialize summarizer
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+    # Stream scum output (appending to the existing combined output)
+    for text in generate_text_stream(scum, tokenizer3, combined_output, max_new_tokens=200, temperature=1.7):
+        scum_output = text
+        combined_output = text
+        yield combined_output
 
-    # Generate the final summary
-    summary = summarize(final_output, summarizer, max_length=500, min_length=150)
+    # Stream gospel output (appending to the existing combined output)
+    for text in generate_text_stream(gospel, tokenizer2, prompt, max_new_tokens=200, temperature=1.0):
+        gospel_output = text
+        combined_output = text
+        yield combined_output
 
-    # Display the final summary
-    yield summary
+    # --- Phase 2: Generate and Stream Summary (Replacing Combined Output) ---
+    final_output_for_summary = scum_output + " " + gospel_output # Use scum and gospel only for summarization
+    for text in generate_summary_stream(summarizer_model, summarizer_tokenizer, final_output_for_summary, max_length=500, min_length=150):
+        yield text
 
 # Gradio interface
 iface = gr.Interface(fn=reply, inputs="text", outputs="text")

requirements.txt

@@ -2,4 +2,4 @@ huggingface_hub==0.25.2
 gradio
 transformers
 bs4 
-torch
+torch