Update app.py

app.py

@@ -4,27 +4,43 @@ os.system("pip install streamlit")
 os.system("pip install torch torchvision")
 
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, GPT2Tokenizer, GPT2LMHeadModel
 
-# Load the GPT-2 model
-generator = pipeline('text-generation', model='gpt2')
+# Load the GPT-2 model and tokenizer
+model_name = 'gpt2'
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
+# Create chatbot pipeline
+chatbot = pipeline('text-generation', model=model, tokenizer=tokenizer)
+
+# Create translation pipeline
+translation = pipeline('translation', model=model, tokenizer=tokenizer)
 
 # Streamlit app
 def main():
-    st.title("GPT-2 Chatbot")
+    st.title("GPT-2 Chatbot & Translation")
 
     # Style the user input text area
     st.markdown('<style>div.Widget.row-widget.stRadio > div{flex-direction:row;}</style>', unsafe_allow_html=True)
     user_input = st.text_area("You:", "Hello, I'm a language model", height=100)
 
+    task = st.radio("Choose Task:", ("Chatbot", "Translation"))
+
     if st.button("Generate"):
-        generated_responses = generator(user_input, max_length=200, num_return_sequences=2)
-                                    
-        for i, response in enumerate(generated_responses):
-            # Style the response box
+        if task == "Chatbot":
+            generated_response = chatbot(user_input, max_length=30, num_return_sequences=1)[0]['generated_text']
+            st.markdown(
+                f'<div style="border: 1px solid #ccc; border-radius: 5px; padding: 10px; margin-bottom: 10px;">'
+                f'Chatbot Response: {generated_response}'
+                '</div>',
+                unsafe_allow_html=True
+            )
+        elif task == "Translation":
+            translated_text = translation(user_input, max_length=30)[0]['translation_text']
             st.markdown(
                 f'<div style="border: 1px solid #ccc; border-radius: 5px; padding: 10px; margin-bottom: 10px;">'
-                f'Response {i + 1}: {response["generated_text"]}'
+                f'Translation: {translated_text}'
                 '</div>',
                 unsafe_allow_html=True
             )