Upload app.py

app.py

@@ -1,28 +1,33 @@
 import streamlit as st
 from transformers import AutoModelForSeq2SeqLM, T5Tokenizer
 
-def generate_response(input_prompt):
-    fine_tuned_model_path = "Robin246/inxai"  # Update with your Hugging Face model path
-
-    # Load the fine-tuned model and tokenizer
-    fine_tuned_model = AutoModelForSeq2SeqLM.from_pretrained(fine_tuned_model_path)
-    tokenizer = T5Tokenizer.from_pretrained(fine_tuned_model_path)
+def generate_response(input_prompt, model_path):
+    if model_path == 'google/flan-t5-small':
+        model_name = 'Google Flan T5'
+    elif model_path == 'MBZUAI/LaMini-Flan-T5-77M':
+        model_name = 'Lamini Flan T5'
+    else:
+        model_name = 'INXAI'
+
+    fine_tuned_model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+    if model_path == 'MBZUAI/LaMini-Flan-T5-77M':
+        tokenizer = T5Tokenizer.from_pretrained('t5-base') 
+    else:
+        tokenizer = T5Tokenizer.from_pretrained(model_path) 
 
     input_text = f"Input prompt: {input_prompt}"
-
     input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=64, padding="max_length", truncation=True)
-    output_ids = fine_tuned_model.generate(input_ids, max_length=256, num_return_sequences=1, num_beams=1, early_stopping=True)
+    output_ids = fine_tuned_model.generate(input_ids, max_length=256, num_return_sequences=1, num_beams=2, early_stopping=True)
     generated_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-    return generated_output
+    return generated_output, model_name
 
 def main():
-    st.title('INXAI LLM model')
-    user_input = st.text_input('Input Prompt:', '')
-
-    if st.button('Generate'):
-        response = generate_response(user_input)
-        st.write('Generated Reply:', response)
-
-if __name__ == '__main__':
+    st.title("INXAI LLM Model")
+    model_selection = st.selectbox("Choose a model", ["google/flan-t5-small", "MBZUAI/LaMini-Flan-T5-77M", "Robin246/inxai"])
+    input_prompt = st.text_input("Enter input text")
+    if st.button("Generate"):
+        reply, model_name = generate_response(input_prompt, model_selection)
+        st.write(f"Generated Reply : {reply}")
+
+if __name__ == "__main__":
     main()