Update app.py

app.py

@@ -1,31 +1,14 @@
-import streamlit as st
 import torch
-from peft import AutoPeftModelForCausalLM
-from transformers import AutoTokenizer, TextStreamer
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 
-# Load LoRA fine-tuned model and tokenizer
-model_path = "lora_model"  # Your model folder path
-load_in_4bit = True  # Whether to load in 4-bit precision
-
-# Load the model
-@st.cache_resource
-def load_model():
-    model = AutoPeftModelForCausalLM.from_pretrained(
-        model_path,
-        torch_dtype=torch.float16 if not load_in_4bit else torch.float32,
-        load_in_4bit=load_in_4bit,
-        device_map="auto"
-    )
-    model.eval()
-    return model
+# Load model from Hugging Face Hub
+base_model = AutoModelForCausalLM.from_pretrained("unsloth/llama-3.2-3b-instruct-bnb-4bit")
+model = PeftModel.from_pretrained(base_model, "ayush0504/Fine-Tunned-GPT")
+model.eval()
 
 # Load tokenizer
-@st.cache_resource
-def load_tokenizer():
-    return AutoTokenizer.from_pretrained(model_path)
-
-model = load_model()
-tokenizer = load_tokenizer()
+tokenizer = AutoTokenizer.from_pretrained("ayush0504/Fine-Tunned-GPT")
 
 def generate_response(question):
     messages = [{"role": "user", "content": question}]
@@ -48,15 +31,11 @@ def generate_response(question):
     
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
-# Streamlit UI
-st.title("Indian Penal Code AI Assistant")
-
-question = st.text_area("Ask a legal question:")
-if st.button("Generate Response"):
+# Example usage
+if __name__ == "__main__":
+    question = input("Ask a legal question: ")
     if question.strip():
-        with st.spinner("Generating response..."):
-            answer = generate_response(question)
-        st.subheader("Answer:")
-        st.write(answer)
+        answer = generate_response(question)
+        print("\nAnswer:", answer)
     else:
-        st.warning("Please enter a question.")
+        print("Please enter a valid question.")