Add Streamlit app and requirements

app.py

@@ -1,34 +1,49 @@
 import streamlit as st
-from transformers import AutoModelForTokenClassification, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import torch
+import bitsandbytes as bnb  # Required for 4-bit quantization
 
-# Load the model and tokenizer from Hugging Face
-model_name = "fajjos/Keyword_v1"  # Replace with the actual model name
+# Load the tokenizer and the quantized LLaMA model
+model_name = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForTokenClassification.from_pretrained(model_name)
+
+# Load the quantized LLaMA model in 4-bit precision
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    load_in_4bit=True,    # Enable 4-bit quantization
+    device_map="auto"     # Automatically assigns to CPU/GPU
+)
+
+# Enable native 2x faster inference (if applicable, ensure this feature works)
+# FastLanguageModel.for_inference(model)  # Uncomment this if FastLanguageModel is available for your model
 
 # Streamlit interface
-st.title("Keyword Extractor")
+st.title("Keyword Extractor using LLaMA 4-bit Model")
+
+# Text input area for user input
 user_input = st.text_area("Enter text for keyword extraction")
 
 if user_input:
-    # Tokenize the input
-    inputs = tokenizer(user_input, return_tensors="pt")
-    
-    # Get model predictions
+    # Prepare the prompt for keyword extraction
+    prompt_template = (
+        "Extract keywords and variables from the prompt:\n"
+        "{}\n"
+    )
+    alpaca_prompt = prompt_template.format(user_input)
+
+    # Tokenize the input text
+    inputs = tokenizer([alpaca_prompt], return_tensors="pt").to("cuda")
+
+    # Set up the text streamer to display the generated text as it streams
+    text_streamer = TextStreamer(tokenizer)
+
+    # Generate keywords and extract variables
     with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Process the predictions (this will depend on your specific model output)
-    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
-    predictions = torch.argmax(outputs.logits, dim=2)
-
-    # Display extracted keywords
-    st.write("Extracted Keywords:")
-    for token, pred in zip(tokens, predictions[0]):
-        if pred == 1:  # Assuming label '1' corresponds to a keyword
-            st.write(token)
-
-# # Add a slider for interaction (example)
-# x = st.slider('Select a value')
-# st.write(f"{x} squared is {x * x}")
+        output = model.generate(**inputs, streamer=text_streamer, max_new_tokens=128)
+
+    # Decode the output tokens to get the generated text
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Display the result in the Streamlit app
+    st.write("Extracted Keywords and Variables:")
+    st.write(generated_text)

requirements.txt

@@ -1,3 +1,5 @@
 transformers
+bitsandbytes
+sentencepiece
 torch
 streamlit