Fix Gemma artifact generation

Signed-off-by: mariecwhite <[email protected]>

common_benchmark_suite/openxla/benchmark/models/jax/gemma/gemma_pipeline.py

@@ -6,6 +6,7 @@
 
 import jax.numpy as jnp
 
+from pathlib import Path
 from transformers import AutoTokenizer, GemmaTokenizer, FlaxPreTrainedModel, FlaxGemmaForCausalLM, GenerationConfig
 from typing import Any, List, Tuple
 
@@ -49,7 +50,8 @@ def __init__(
     self.tokenizer = AutoTokenizer.from_pretrained(
         model_name,
         model_max_length=self.seq_len,
-        padding_side="left",
+        max_length=self.seq_len,
+        padding="max_length",
     )
     self.tokenizer.pad_token = self.tokenizer.eos_token
     self.tokenization_kwargs = {
@@ -63,18 +65,20 @@ def __init__(
         use_cache=True)
 
   def generate_default_inputs(self) -> str:
-    return "Once upon a time"
+    input_file = Path(__file__).parent / "summarization_prompt.txt"
+    with open(input_file, 'r', encoding='utf-8') as file:
+      text = file.read()
+    return text
 
   def preprocess(self, input_text: str) -> Tuple[Any,]:
     batch_input_text = [input_text] * self.batch_size
     inputs = self.tokenizer(text=batch_input_text, **self.tokenization_kwargs)
     return (inputs["input_ids"],)
 
   def forward(self, input_text: Any) -> Any:
-    output = self.model.generate(input_text,
-                                 params=self.params,
-                                 generation_config=self.generation_config)
-    print(f"output: {output}")
+    return self.model.generate(input_text,
+                               params=self.params,
+                               generation_config=self.generation_config)
 
   def postprocess(self, output: Any) -> List[str]:
     return self.tokenizer.batch_decode(output, skip_special_tokens=True)

common_benchmark_suite/openxla/benchmark/models/jax/gemma/summarization_prompt.txt

@@ -0,0 +1,42 @@
+Title: The Rise and Fall of the Second Galactic Empire: A Comprehensive Socioeconomic and Political Analysis.
+
+Abstract: This historical analysis delves into the intricate factors that contributed to the meteoric ascent and subsequent collapse of the Second Galactic Empire, a dominant interstellar power that spanned over three millennia. Drawing upon a wealth of primary sources, including imperial records, personal journals, and archaeological artifacts, this study examines the empire's multifaceted socioeconomic structure, intricate political machinations, technological advancements, and cultural shifts.
+
+The research identifies key catalysts for the empire's initial expansion, such as groundbreaking innovations in hyperspace travel, astute diplomatic alliances, and a centralized bureaucratic system that facilitated efficient resource allocation and governance. However, the study also reveals the seeds of the empire's decline, including growing social inequality, escalating corruption within the ruling elite, over-extension of military resources, and a gradual erosion of public trust in imperial institutions.
+
+Through a meticulous examination of these interconnected elements, this analysis aims to shed light on the complex interplay of forces that shaped the trajectory of the Second Galactic Empire. The findings of this research not only offer valuable insights into the dynamics of past civilizations but also hold potential implications for contemporary societies grappling with similar challenges of growth, governance, and sustainability.
+
+Introduction: The Second Galactic Empire stands as a monumental chapter in the annals of interstellar history. Its rise to power was swift and awe-inspiring, propelled by a confluence of technological prowess, political acumen, and a seemingly inexhaustible thirst for expansion. However, the empire's reign was not without its shadows. Beneath the veneer of prosperity and order lurked simmering discontent, social unrest, and a growing sense of disillusionment among the empire's diverse populace.
+
+This study seeks to unravel the intricate tapestry of events, decisions, and societal trends that ultimately led to the empire's demise. By examining the empire's socioeconomic structure, political landscape, technological achievements, and cultural evolution, we aim to uncover the underlying causes of its rise and fall.
+
+Methods:
+
+Primary Source Analysis: This research draws extensively from a wide range of primary sources, including:
+- Imperial Records: Official decrees, edicts, census data, and administrative reports from various sectors of the empire, including the Core Worlds, the Outer Rim, and newly conquered territories. These records provide valuable insights into the empire's administrative structure, economic policies, and military strategies.
+- Personal Journals: Diaries, letters, and memoirs of prominent figures, such as Emperor Palpatine, Grand Moff Tarkin, and Princess Leia Organa, as well as lesser-known individuals from various social strata. These personal accounts offer a glimpse into the motivations, fears, and aspirations of those who lived under the empire's rule.
+- Archaeological Artifacts: Recovered structures, such as the ruins of Imperial palaces and military bases, as well as tools, weapons, artwork, and everyday objects. These artifacts provide tangible evidence of the empire's technological capabilities, artistic achievements, and daily life.
+- Comparative Historical Analysis: We employ a comparative approach, drawing parallels between the Second Galactic Empire and other historical empires on Earth and beyond, such as the Roman Empire, the Mongol Empire, and the British Empire. This allows us to identify recurring patterns, such as the challenges of maintaining a vast and diverse empire, as well as unique aspects of the Galactic Empire's development, such as the role of advanced technology and interstellar travel.
+- Quantitative and Qualitative Analysis: We utilize both quantitative methods, such as statistical analysis of economic data and demographic trends, and qualitative methods, such as textual analysis of historical documents and discourse analysis of political speeches. This multi-faceted approach allows us to gain a nuanced understanding of the complex interplay of factors that shaped the empire's trajectory.
+
+Key Research Questions:
+
+- What were the primary factors that propelled the Second Galactic Empire's rapid expansion and consolidation of power?
+- How did the empire's socioeconomic structure evolve over time, and what were the implications of this evolution for social stability and political legitimacy?
+- What role did technological advancements play in both the empire's successes and its eventual vulnerabilities?
+- How did the empire's political system function, and what were the key decision-making processes that shaped its policies and actions?
+- What were the underlying social, economic, and political grievances that fueled discontent and ultimately contributed to the empire's collapse?
+- What lessons can be learned from the rise and fall of the Second Galactic Empire that might be applicable to contemporary societies facing similar challenges?
+
+Prompt for the LLM:
+
+Please generate a comprehensive summary (approximately 300-400 words) of this historical analysis.
+
+Focus on:
+- Key research questions and objectives of the study.
+- Primary sources utilized in the research.
+- Major findings regarding the factors contributing to the empire's rise and fall.
+- Implications of the study for understanding historical and contemporary civilizations.
+- Any potential limitations or areas for further research identified in the analysis.
+
+Add conclusion.