Update app.py

app.py

@@ -1,49 +1,56 @@
+import os
 import streamlit as st
-from huggingface_hub import InferenceClient
-
-global base_url
-global token
-
-# Initialize the Hugging Face Inference Client
-client = InferenceClient(
-    base_url,
-    token,
-)
-
-# Function to generate project tasks using LLaMA 3 8B
-def generate_project_tasks(domain):
-    prompt = f"Create two simple project tasks for an internship in {domain}."
-    tasks = ""
-    for message in client.chat_completion(
-        messages=[{"role": "user", "content": prompt}],
-        max_tokens=500,
-        stream=True,
-    ):
-        tasks += message.choices[0].delta.content
-    return tasks
-
-# Streamlit app layout
-st.title("AI Internship Task Generator")
-st.write("Enter the internship domain to generate two project tasks:")
-
-# Input field for internship domain
-domain = st.text_input("Internship Domain", "")
-
-# Button to generate tasks
-if st.button("Generate Tasks"):
-    if domain:
-        # Generate tasks and display them
-        tasks = generate_project_tasks(domain)
-        st.write("### Generated Project Tasks")
-        st.write(tasks)
-
-        # Save tasks to a text file
-        file_name = f"{domain}_project_tasks.txt"
-        with open(file_name, "w") as file:
-            file.write(tasks)
-
-        st.write(f"Tasks have been saved to {file_name}.")
-
-    else:
-        st.error("Please enter a domain.")
-
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Load model and tokenizer
+def load_model(model_name):
+    global model
+    global tokenizer
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Generate response based on input text
+def generate_response(input_text):
+    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    output_ids = model.generate(
+        input_ids,
+        no_repeat_ngram_size=2,
+        max_new_tokens=200,
+        num_beams=2,
+        eos_token_id=tokenizer.eos_token_id
+    )
+    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+# Streamlit app
+def main():
+    # Fetch the model name and token from environment variables
+    model_name = os.getenv("MODEL_NAME", "meta-llama/Meta-Llama-3-8B")
+    hf_token = os.getenv("HF_TOKEN")
+
+    # Verify if the token is set
+    if not hf_token:
+        st.error("Hugging Face token not found. Please set it in your environment variables.")
+        return
+
+    # Load the model
+    load_model(model_name)
+
+    st.title("AI Internship Task Generator")
+    st.write("Enter an internship domain to generate tasks.")
+
+    # Input field for the internship domain
+    input_text = st.text_input("Enter Internship Domain:")
+
+    # Button to generate tasks
+    if st.button("Generate Tasks"):
+        if input_text:
+            with st.spinner("Generating tasks..."):
+                # Generate response using the model
+                response = generate_response(input_text)
+                st.write("**Generated Tasks:**")
+                st.write(response)
+        else:
+            st.warning("Please enter an internship domain.")
+
+if __name__ == "__main__":
+    main()