Amelia-James commited on
Commit
f635cd7
·
verified ·
1 Parent(s): ff617a6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +80 -0
app.py ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import numpy as np
3
+ import torch
4
+ from transformers import AutoTokenizer, AutoModel, RagTokenizer, RagRetriever, RagSequenceForGeneration
5
+ from pymilvus import connections, Collection, CollectionSchema, FieldSchema, DataType
6
+ from dotenv import load_dotenv
7
+ import os
8
+
9
+ # Load environment variables
10
+ load_dotenv()
11
+ GROQ_API_KEY = os.getenv('GROQ_API_KEY')
12
+
13
+ # Initialize Milvus connection
14
+ connections.connect("default", host="localhost", port="19530")
15
+
16
+ # Define Milvus schema and collection
17
+ fields = [
18
+ FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
19
+ FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768) # Adjust the dimension based on your model
20
+ ]
21
+ schema = CollectionSchema(fields, "User Data Collection")
22
+ collection = Collection(name="user_data", schema=schema)
23
+
24
+ # Load Hugging Face models
25
+ tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
26
+ model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
27
+ tokenizer_rag = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
28
+ retriever = RagRetriever.from_pretrained("facebook/rag-sequence-nq", index_name="custom")
29
+ model_rag = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
30
+
31
+ # Define functions
32
+ def generate_embedding(text):
33
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
34
+ with torch.no_grad():
35
+ outputs = model(**inputs)
36
+ return outputs.last_hidden_state.mean(dim=1).numpy().tolist()[0]
37
+
38
+ def insert_data(user_id, embedding):
39
+ collection.insert([user_id, embedding])
40
+
41
+ def retrieve_relevant_data(query):
42
+ query_embedding = generate_embedding(query)
43
+ search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
44
+ results = collection.search(query_embedding, "embedding", search_params)
45
+ return results
46
+
47
+ def generate_cv(job_description, company_profile=None):
48
+ query = job_description
49
+ if company_profile:
50
+ query += f" Company profile: {company_profile}"
51
+
52
+ relevant_data = retrieve_relevant_data(query)
53
+ context = " ".join([data.text for data in relevant_data])
54
+
55
+ inputs = tokenizer_rag(query, return_tensors="pt")
56
+ context_inputs = tokenizer_rag(context, return_tensors="pt")
57
+ outputs = model_rag.generate(input_ids=inputs['input_ids'], context_input_ids=context_inputs['input_ids'])
58
+ return tokenizer_rag.decode(outputs[0], skip_special_tokens=True)
59
+
60
+ # Streamlit UI
61
+ st.title("Custom CV Generator")
62
+
63
+ st.sidebar.header("Input Data")
64
+ skills = st.sidebar.text_input("Enter your skills")
65
+ experience = st.sidebar.text_input("Enter your experience")
66
+ education = st.sidebar.text_input("Enter your education")
67
+ job_description = st.sidebar.text_area("Enter job description")
68
+ company_profile = st.sidebar.text_area("Enter company profile (optional)")
69
+
70
+ if st.sidebar.button("Generate CV"):
71
+ # Insert user data (assuming single user for simplicity)
72
+ user_data = f"Skills: {skills}. Experience: {experience}. Education: {education}."
73
+ user_id = 1 # Example user ID
74
+ user_embedding = generate_embedding(user_data)
75
+ insert_data(user_id, user_embedding)
76
+
77
+ # Generate CV
78
+ cv_text = generate_cv(job_description, company_profile)
79
+ st.write("Your Tailored CV:")
80
+ st.write(cv_text)