Update app.py

app.py

@@ -1,4 +1,6 @@
 import streamlit as st
+import torch
+import re
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
 # Load the fine-tuned model and tokenizer
@@ -6,6 +8,39 @@ model_name = "rohangbs/fine-tuned-gpt2"
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 model = GPT2LMHeadModel.from_pretrained(model_name)
 
+# Ensure the model is on the correct device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Function to generate a response
+def chat_with_model(input_prompt, max_length=200):
+    model.eval()
+    
+    # Format the input prompt with special tokens
+    prompt = f"<|startoftext|>[WP] {input_prompt}\n[RESPONSE]"
+    
+    # Tokenize and encode the prompt, and send to the device
+    generated = torch.tensor(tokenizer.encode(prompt)).unsqueeze(0).to(device)
+    
+    # Generate a response
+    sample_outputs = model.generate(
+        generated,
+        do_sample=True,
+        top_k=50,
+        max_length=max_length,
+        top_p=0.95,
+        num_return_sequences=1,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    
+    # Decode the response and clean it up
+    response_text = tokenizer.decode(sample_outputs[0], skip_special_tokens=True)
+    wp_responses = re.split(r"\[WP\].*?\n|\[RESPONSE\]", response_text)[1:]
+    clean_responses = [response.strip() for response in wp_responses if response.strip()]
+    
+    # Return the first valid response
+    return clean_responses[0] if clean_responses else "I couldn't generate a response."
+
 # Streamlit UI
 st.title("Chatbot For Company Details")
 st.write("A GPT-2 model fine-tuned for Company dataset.")
@@ -16,14 +51,10 @@ prompt = st.text_area("Ask your question:", height=150)
 if st.button("Send"):
     if prompt.strip():
         with st.spinner("Generating..."):
-            # Generate text
-            inputs = tokenizer.encode(prompt, return_tensors="pt")
-            outputs = model.generate(inputs, max_length=200, num_return_sequences=1)
-            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Display the response
-        st.subheader("Generated Response:")
-        st.write(response)
+            # Generate and display the response
+            response = chat_with_model(prompt)
+            st.subheader("Generated Response:")
+            st.write(response)
     else:
         st.warning("Please enter a prompt.")