Update app.py

app.py

@@ -1,25 +1,50 @@
-from transformers import MarianMTModel, MarianTokenizer
-
-# Load the model and tokenizer for English-to-French translation
-model_name = "Helsinki-NLP/opus-mt-en-fr"  # English to French translation model
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
-
-def translate_text(input_text):
-    # Tokenize the input text
-    input_text = ">>en<< " + input_text  # Prefix the input text with the source language code (en for English)
-    inputs = tokenizer(input_text, return_tensors="pt")
-
-    # Perform translation
-    with torch.no_grad():
-        outputs = model.generate(**inputs)
-
-    # Decode the translated text
-    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return translated_text
-
-# Example usage
-input_text = "Hello, how are you?"
-translated_text = translate_text(input_text)
-print("English:", input_text)
-print("French:", translated_text)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+import torch
+
+title = "🤖AI ChatBot"
+description = "A State-of-the-Art Large-scale Pretrained Response generation model (gpt-neo-1.3B)"
+examples = [["How are you?"]]
+
+# Use the better model and tokenizer
+model_name = "EleutherAI/gpt-neo-1.3B"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+def predict(input_text, history=None):
+    if history is None:
+        history = []
+
+    # Tokenize the new input sentence
+    new_user_input_ids = tokenizer.encode(
+        input_text + tokenizer.eos_token, return_tensors="pt"
+    )
+
+    # Append the new user input tokens to the chat history
+    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
+
+    # Generate a response using batch processing
+    generated_ids = model.generate(
+        bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id
+    )
+
+    # Convert the generated response tokens to text
+    response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+    # Split the responses into lines
+    response = response.split("\n")
+
+    # Convert to tuples of list
+    response = [(response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)]
+
+    return response, generated_ids.tolist()
+
+gr.Interface(
+    fn=predict,
+    title=title,
+    description=description,
+    examples=examples,
+    inputs=["text", "state"],
+    outputs=["chatbot", "state"],
+    theme="finlaymacklon/boxy_violet",
+).launch()