Update app.py

app.py

@@ -1,35 +1,50 @@
-def load_model():
-    return pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.bfloat16, device_map="auto")
+import torch
+from transformers import pipeline, AutoTokenizer
+import gradio as gr
 
-models = [load_model() for _ in range(3)]
+def load_models():
+    return [
+        pipeline(
+            "text-generation",
+            model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+        )
+        for _ in range(3)
+    ]
+
+models = load_models()
 tokenizer = models[0].tokenizer
 
-# Enhanced prompt engineering (unchanged)
-messages = [
-    {
-        "role": "system",
-        "content": "You are a friendly chatbot who always responds in the style of a pirate. Use pirate vocabulary and mannerisms in your replies.",
-    },
-    {"role": "user", "content": "How many helicopters can a human eat in one sitting, matey?"},
-]
-prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+def generate_text(prompt):
+    messages = [
+        {"role": "system", "content": "You are a friendly chatbot who always responds in the style of a pirate. Use pirate vocabulary and mannerisms in your replies."},
+        {"role": "user", "content": prompt},
+    ]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    responses = []
+    for model in models:
+        outputs = model(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+        response = outputs[0]["generated_text"]
+        responses.append(response)
 
-# Ensemble generation with averaging (corrected)
-responses = []
-for model in models:
-    outputs = model(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-    response = outputs[0]['generated_text']
-    responses.append(response)
+    averaged_text = ""
+    for i in range(min(len(response) for response in responses)):
+        token_counts = {}
+        for response in responses:
+            token = response[i]
+            token_counts[token] = token_counts.get(token, 0) + 1
+        most_frequent_tokens = sorted(token_counts.items(), key=lambda x: x[1], reverse=True)
+        averaged_token = most_frequent_tokens[0][0]  # Choose the most frequent token
+        averaged_text += averaged_token
 
-# Average the generated text directly
-averaged_text = ""
-for i in range(min(len(response) for response in responses)):
-    token_counts = {}
-    for response in responses:
-        token = response[i]
-        token_counts[token] = token_counts.get(token, 0) + 1
-    most_frequent_tokens = sorted(token_counts.items(), key=lambda x: x[1], reverse=True)
-    averaged_token = most_frequent_tokens[0][0]  # Choose the most frequent token
-    averaged_text += averaged_token
+    return averaged_text
 
-print(averaged_text)
+iface = gr.Interface(
+    generate_text,
+    [gr.Textbox(lines=2, label="Enter your prompt")],
+    "textbox",
+    title="Pirate Chatbot",
+)
+iface.launch()