major code adjustments to hopefully reduce latency

app.py

@@ -1,50 +1,65 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import concurrent.futures
 
-# Set Device to CPU
-device = torch.device('cpu')
+# Set Device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load the Models
-# Model 1: GPT-2 Medium
-tokenizer1 = AutoTokenizer.from_pretrained('gpt2-medium')
-model1 = AutoModelForCausalLM.from_pretrained('gpt2-medium')
+# Load Models
+# Model 1: Bloom 560M
+tokenizer1 = AutoTokenizer.from_pretrained('bigscience/bloom-560m')
+model1 = AutoModelForCausalLM.from_pretrained('bigscience/bloom-560m', torch_dtype=torch.float16)
 model1.to(device)
 
-# Model 2: GPT-Neo 125M
-tokenizer2 = AutoTokenizer.from_pretrained('EleutherAI/gpt-neo-125M')
-model2 = AutoModelForCausalLM.from_pretrained('EleutherAI/gpt-neo-125M')
+# Model 2: GPT-Neo 1.3B
+tokenizer2 = AutoTokenizer.from_pretrained('EleutherAI/gpt-neo-1.3B')
+model2 = AutoModelForCausalLM.from_pretrained('EleutherAI/gpt-neo-1.3B', torch_dtype=torch.float16)
 model2.to(device)
 
-# Define Text Generation Functions
-def generate_text_model1(prompt, temperature, top_p):
-    inputs = tokenizer1(prompt, return_tensors='pt').to(device)
+# Define Functions with Improved Parameters
+def generate_text_model1(prompt):
+    inputs = tokenizer1.encode(prompt, return_tensors='pt').to(device)
     with torch.no_grad():
         outputs = model1.generate(
-            **inputs,
-            max_new_tokens=30,
+            inputs,
+            max_length=50,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
             do_sample=True,
             top_k=50,
-            top_p=top_p,
-            temperature=temperature
+            top_p=0.95,
+            temperature=0.8
         )
     text = tokenizer1.decode(outputs[0], skip_special_tokens=True)
     return text
 
-def generate_text_model2(prompt, temperature, top_p):
-    inputs = tokenizer2(prompt, return_tensors='pt').to(device)
+def generate_text_model2(prompt):
+    inputs = tokenizer2.encode(prompt, return_tensors='pt').to(device)
     with torch.no_grad():
         outputs = model2.generate(
-            **inputs,
-            max_new_tokens=30,
+            inputs,
+            max_length=50,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
             do_sample=True,
             top_k=50,
-            top_p=top_p,
-            temperature=temperature
+            top_p=0.95,
+            temperature=0.8
         )
     text = tokenizer2.decode(outputs[0], skip_special_tokens=True)
     return text
 
+# Use ThreadPoolExecutor to Process in Parallel
+def compare_models(prompt):
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future1 = executor.submit(generate_text_model1, prompt)
+        future2 = executor.submit(generate_text_model2, prompt)
+        output1 = future1.result()
+        output2 = future2.result()
+    return output1, output2
+
+
 def compare_models(prompt, temperature, top_p):
     output1 = generate_text_model1(prompt, temperature, top_p)
     output2 = generate_text_model2(prompt, temperature, top_p)
@@ -55,21 +70,23 @@ def compare_models(prompt, temperature, top_p):
 
     return output1_with_params, output2_with_params
 
+# Use ThreadPoolExecutor to Process in Parallel
+def compare_models(prompt):
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future1 = executor.submit(generate_text_model1, prompt)
+        future2 = executor.submit(generate_text_model2, prompt)
+        output1 = future1.result()
+        output2 = future2.result()
+    return output1, output2
+
 # Create Gradio Interface
 iface = gr.Interface(
     fn=compare_models,
-    inputs=[
-        gr.Textbox(lines=2, placeholder='Enter a prompt here...', label='Prompt'),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.8, label='Temperature'),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label='Top-p')
-    ],
-    outputs=[
-        gr.Markdown(label='GPT-2 Medium Output'),
-        gr.Markdown(label='GPT-Neo 125M Output')
-    ],
-    title='Compare Text Generation Models with Adjustable Parameters',
-    description='Enter a prompt and adjust the temperature and top-p parameters to see how they affect the generated text.'
+    inputs=gr.Textbox(lines=2, placeholder='Enter a prompt here...'),
+    outputs=[gr.Textbox(label='Bloom 560M Output'), gr.Textbox(label='GPT-Neo 1.3B Output')],
+    title='Compare Text Generation Models',
+    description='Enter a prompt and see how two different models generate text.'
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# Launch Interface
+iface.launch()