Update app.py

app.py

@@ -82,12 +82,16 @@ def hf_generate(model_name, prompt, max_new_tokens=256, temperature=0.2):
     client = InferenceClient(model=model_name, token=HF_TOKEN)
     start = time.time()
     try:
-        output = client.text_generation(prompt, max_new_tokens=max_new_tokens, temperature=temperature)
+        if "flan" in model_name or "t5" in model_name:
+            output = client.text2text_generation(prompt, max_new_tokens=max_new_tokens)
+        else:
+            output = client.text_generation(prompt, max_new_tokens=max_new_tokens, temperature=temperature)
         latency = time.time() - start
         return output.strip(), latency
     except Exception as e:
         return f"ERROR: {e}", time.time() - start
 
+
 # ---------------- Benchmark Function ---------------- #
 def benchmark(config_text, dataset_text, task):
     cfg = yaml.safe_load(config_text)