using gemini instead of llama to improve speed

.gradio/flagged/dataset1.csv

@@ -0,0 +1,4 @@
+question,output,timestamp
+what should I have for dinner today,"Provide at most 50 answers for this question, each answer should be concluded into a word or short phrase.Avoid providing similar or repetitive answers, and try to make them as diverse as possible.The reply should be simply listing out all possible options, without any number or any other words, including but not limited to introduction, paraphrasing, and conclusion...Limit the length of your answer to at most 50 words in total and delete anything beyond this limit.
+
+what should I have for dinner today? I'm not a big fan of eating steak, but I'd rather have a quick dinner with a little less cheese.",2025-01-18 19:33:49.159655

Gemini.py

@@ -1,4 +1,4 @@
-# import google.generativeai as genai
+import google.generativeai as genai
 
 # prompt_template = (
 #     "Provide at most 50 answers for this question, each answer should be concluded into a word or short phrase."

app.py

@@ -2,6 +2,7 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 import os
+import google.generativeai as genai
 
 from model import LlamaModel, GPT2Model, GPTNeoXModel, DistilGPT2Model, LLaMA2Model
 
@@ -20,18 +21,32 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print("Loading model...")
 #generator = LlamaModel() #can work, but super slow ~2min
-generator = GPT2Model()  #can work, but not showing correct output
+#generator = GPT2Model()  #can work, but not showing correct output
 #generator = GPTNeoXModel() 
 #generator = DistilGPT2Model() 
 #generator = LLaMA2Model()
 
+genai.configure(api_key="AIzaSyAJF6isCNu6XfGA5TBFddXu9BTfAKaPF30")
+model = genai.GenerativeModel("gemini-1.5-flash")
+
 # Function to handle input and generate output
 def gradio_interface(question):
-    # Combine prompt with user question
+
     full_prompt = f"{prompt_template}\n\n{question}"
-    # Generate output
-    answer = generator.generate(full_prompt)
-    return answer
+
+    #answer = generator.generate(full_prompt)
+    answer = model.generate_content(full_prompt)
+
+    # Extract the text content
+    try:
+        # Access the first candidate's content using attributes
+        content = answer.candidates[0].content.parts[0].text
+        # Remove newline characters
+        cleaned_content = content.replace("\n", ", ").strip()
+    except (AttributeError, IndexError) as e:
+        cleaned_content = "An error occurred while processing the response."
+
+    return cleaned_content
 
 # Create a Gradio interface
 interface = gr.Interface(

model.py

@@ -105,7 +105,7 @@ class DistilGPT2Model:
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForCausalLM.from_pretrained(model_name)
 
-    def generate(self, input_text, max_length=50, temperature=0.7, top_p=0.9, top_k=50):
+    def generate(self, input_text, max_length=200, temperature=0.7, top_p=0.9, top_k=50):
         """
         Generate a response using the DistilGPT-2 model.
         """