Update app.py

app.py

@@ -1,19 +1,12 @@
 import streamlit as st
 import requests
 from bs4 import BeautifulSoup
-#from transformers import AutoTokenizer, AutoModelForCausalLM # Removed for GGUF
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import json
+import torch
 import re
-from llama_cpp import Llama # For GGUF
-import os # For checking GGUF file
 
-# --- Constants ---
-GGUF_MODEL_PATH = "gemma-2-2b-it.Q4_K_M.gguf"  # Or your GGUF file path
-if not os.path.exists(GGUF_MODEL_PATH):
-  st.error(f"Error: GGUF model not found at: {GGUF_MODEL_PATH}. Please download it and place it next to the app.")
-  st.stop()
-
-# --- Scraping Function (No changes needed) ---
+# Function to scrape a URL using Beautiful Soup
 def scrape_url(url):
     try:
         headers = {
@@ -32,26 +25,18 @@ def scrape_url(url):
         return None
 
 
-# --- Load GGUF model with llama-cpp ---
 @st.cache_resource
 def load_model():
-    try:
-        llm = Llama(model_path=GGUF_MODEL_PATH, n_gpu_layers=20) # Change n_gpu_layers if needed
-        return llm
-    except Exception as e:
-         st.error(f"Error loading model: {e}")
-         return None
-
+    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it", trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it",
+                                                torch_dtype=torch.bfloat16,
+                                                low_cpu_mem_usage=True,
+                                                device_map="auto",
+                                                trust_remote_code=True)
+    return tokenizer, model
 
-model = load_model()
+tokenizer, model = load_model()
 
-if model is None:
-    st.error("Failed to load the model. Please check the logs for errors.")
-    st.stop()
-
-
-
-# --- Generate JSON Output with Llama-cpp ---
 def generate_json_output(text):
     prompt = f"""You are a web page text scanner. Your task is to carefully review text from a web page.
 
@@ -66,9 +51,10 @@ Answer the following questions:
 You should output your answers strictly in the following JSON format, but do NOT use markdown:
 {{\"brand\": \"<brand>\", \"intent\": \"<intent>\"}}
 """
-    
-    output = model(prompt, max_tokens=256, stop=["\n"], echo=False)
-    response = output['choices'][0]['text']
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=256, return_dict_in_generate=True)
+    generated_tokens = outputs.sequences[:, inputs.input_ids.shape[1]:]
+    response = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
 
     output_json = None
     try:
@@ -86,7 +72,7 @@ You should output your answers strictly in the following JSON format, but do NOT
     
     return response, output_json
 
-# --- Streamlit App ---
+# Streamlit app
 def main():
     st.title("Google Brand and Intent Detection")
     st.write("Google's brand and intent detection reverse engineered from Chrome by [DEJAN AI](https://dejan.ai/).")