Update app.py

app.py

@@ -1,35 +1,31 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load a smaller model and tokenizer
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Replace with any model name from the list
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Function to generate a response from the chatbot
-def chat_with_bot(prompt):
-    # Encode the input prompt
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-
-    # Generate response
-    with torch.no_grad():
-        output = model.generate(input_ids, max_length=150, num_return_sequences=1)
+# Load model and tokenizer
+@st.cache_resource
+def load_model():
+    model_name = "microsoft/Phi-3-mini-4k-instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return model, tokenizer
 
-    # Decode the generated text
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
+def generate_text(prompt, model, tokenizer):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs.input_ids, max_length=300, temperature=0.7, top_k=50, top_p=0.95)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
 
-# Streamlit UI
-st.title("My Boss Glaston")
-st.write("Chatbot: Hello! I'm here to help you. Type your message below.")
+# Streamlit app
+def main():
+    st.title("Instruction-Following Model: Phi-3-Mini-4k-Instruct")
+    st.write("Ask a question or give an instruction to get a response.")
+
+    model, tokenizer = load_model()
 
-# Input box for user messages
-user_input = st.text_input("You:")
+    prompt = st.text_input("Enter your prompt:", "Explain the concept of machine learning in simple terms.")
+    if st.button("Generate Response"):
+        response = generate_text(prompt, model, tokenizer)
+        st.write(response)
 
-if st.button("Send"):
-    if user_input:
-        response = chat_with_bot(user_input)
-        st.text_area("Chatbot:", value=response, height=100, disabled=True)
-    else:
-        st.text_area("Chatbot:", value="Please type a message!", height=100, disabled=True)
+if __name__ == "__main__":
+    main()