Update app.py

app.py

@@ -1,11 +1,9 @@
 import streamlit as st
-from transformers import pipeline
-from scrapegraphai.graphs import SmartScraperGraph
-import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Page config
 st.set_page_config(
-    page_title="Zephyr Chat & Scrape",
+    page_title="Zephyr Chat",
     page_icon="🤖",
     layout="wide"
 )
@@ -14,98 +12,56 @@ st.set_page_config(
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-if "scrape_results" not in st.session_state:
-    st.session_state.scrape_results = None
-
-# Load Zephyr model
+# Load model and tokenizer
 @st.cache_resource
 def load_model():
-    try:
-        return pipeline(
-            "text-generation",
-            model="HuggingFaceH4/zephyr-7b-beta",
-            torch_dtype=torch.float16,
-            device_map="auto",
-            model_kwargs={"load_in_8bit": True}  # Use 8-bit quantization to reduce memory usage
-        )
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None
-
-# Initialize the model
-model = load_model()
-
-# Sidebar for web scraping
-with st.sidebar:
-    st.title("Web Scraping")
-    url = st.text_input("Enter URL to scrape")
-    scrape_prompt = st.text_input("What information do you want to extract?")
-    
-    if st.button("Scrape"):
-        try:
-            # Configure scraper
-            graph_config = {
-                "llm": {
-                    "model": "HuggingFaceH4/zephyr-7b-beta",
-                    "temperature": 0.7,
-                },
-                "verbose": True
-            }
-            
-            # Create scraper instance
-            scraper = SmartScraperGraph(
-                prompt=scrape_prompt,
-                source=url,
-                config=graph_config
-            )
-            
-            # Run scraping
-            st.session_state.scrape_results = scraper.run()
-            st.success("Scraping completed!")
-            
-        except Exception as e:
-            st.error(f"Error during scraping: {str(e)}")
+    model_name = "HuggingFaceH4/zephyr-7b-beta"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return model, tokenizer
 
 # Main chat interface
 st.title("Zephyr Chatbot 🤖")
 
-# Display scraped results if available
-if st.session_state.scrape_results:
-    st.subheader("Scraped Information")
-    st.json(st.session_state.scrape_results)
+try:
+    model, tokenizer = load_model()
     
-# Display chat messages
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Chat input
+    if prompt := st.chat_input("What's on your mind?"):
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        
+        # Generate response
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                # Prepare input
+                input_text = f"User: {prompt}\nAssistant:"
+                inputs = tokenizer(input_text, return_tensors="pt")
+                
+                # Generate response
+                outputs = model.generate(
+                    inputs.input_ids,
+                    max_length=200,
+                    num_return_sequences=1,
+                    temperature=0.7,
+                    pad_token_id=tokenizer.eos_token_id
+                )
+                
+                # Decode and display response
+                response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+                response = response.split("Assistant:")[-1].strip()
+                
+                st.markdown(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+
+except Exception as e:
+    st.error(f"Error: {str(e)}")
+    st.info("Note: This app requires significant computational resources. Consider using a smaller model or upgrading your Space's resources.")
 
-# Chat input
-if prompt := st.chat_input("What's on your mind?"):
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.markdown(prompt)
-    
-    # Generate response
-    with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            # Include scraped content in context if available
-            context = ""
-            if st.session_state.scrape_results:
-                context = f"Scraped information: {str(st.session_state.scrape_results)}\n"
-            
-            full_prompt = f"{context}User: {prompt}\nAssistant:"
-            
-            response = model(
-                full_prompt,
-                max_length=1000,
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.15
-            )[0]["generated_text"]
-            
-            # Clean up response to get only the assistant's reply
-            response = response.split("Assistant:")[-1].strip()
-            
-            st.markdown(response)
-            st.session_state.messages.append({"role": "assistant", "content": response})